DerivePrimaryKey with custom primary key column name (#694)

* `DerivePrimaryKey` with custom primary key column name

* Add test cases [issues]

.github/workflows/rust.yml

@@ -318,7 +318,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
-        path: [86, 249, 262, 319, 324, 352, 356, 471]
+        path: [86, 249, 262, 319, 324, 352, 356, 471, 693]
     steps:
       - uses: actions/checkout@v2
 

issues/693/Cargo.toml

@@ -0,0 +1,22 @@
+[workspace]
+# A separate workspace
+
+[package]
+name = "sea-orm-issues-693"
+version = "0.1.0"
+authors = ["bleuse <raphael.bleuse@univ-grenoble-alpes.fr>"]
+edition = "2021"
+publish = false
+
+[dependencies]
+tokio = { version = "1.14", features = ["full"] }
+anyhow = "1"
+dotenv = "0.15"
+futures-util = "0.3"
+serde = "1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+
+[dependencies.sea-orm]
+path = "../../" # remove this line in your own project
+features = ["runtime-tokio-rustls", "sqlx-mysql", "macros"]
+default-features = false

issues/693/src/container.rs

@@ -0,0 +1,31 @@
+pub mod prelude {
+    pub use super::model::{
+        ActiveModel as ContainerActiveModel, Column as ContainerColumn, Entity as Container,
+        Model as ContainerModel, PrimaryKey as ContainerPrimaryKey, Relation as ContainerRelation,
+    };
+}
+
+pub mod model {
+    use sea_orm::entity::prelude::*;
+
+    #[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+    #[sea_orm(table_name = "container")]
+    pub struct Model {
+        #[sea_orm(primary_key, column_name = "db_id")]
+        pub rust_id: i32,
+    }
+
+    #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+    pub enum Relation {
+        #[sea_orm(has_many = "crate::Content")]
+        Content, // 1(Container) ⇆ n(Content)
+    }
+
+    impl Related<crate::Content> for Entity {
+        fn to() -> RelationDef {
+            Relation::Content.def()
+        }
+    }
+
+    impl ActiveModelBehavior for ActiveModel {}
+}

issues/693/src/content.rs

@@ -0,0 +1,36 @@
+pub mod prelude {
+    pub use super::model::{
+        ActiveModel as ContentActiveModel, Column as ContentColumn, Entity as Content,
+        Model as ContentModel, PrimaryKey as ContentPrimaryKey, Relation as ContentRelation,
+    };
+}
+
+pub mod model {
+    use sea_orm::entity::prelude::*;
+
+    #[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+    #[sea_orm(table_name = "content")]
+    pub struct Model {
+        #[sea_orm(primary_key)]
+        pub id: i32,
+        pub container_id: i32,
+    }
+
+    #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+    pub enum Relation {
+        #[sea_orm(
+            belongs_to = "crate::Container",
+            from = "crate::ContentColumn::ContainerId",
+            to = "crate::ContainerColumn::RustId"
+        )]
+        Container, // 1(Container) ⇆ n(Content)
+    }
+
+    impl Related<crate::Container> for Entity {
+        fn to() -> RelationDef {
+            Relation::Container.def()
+        }
+    }
+
+    impl ActiveModelBehavior for ActiveModel {}
+}

issues/693/src/main.rs

@@ -0,0 +1,18 @@
+mod container;
+mod content;
+
+use container::prelude::*;
+use content::prelude::*;
+use sea_orm::{DbBackend, EntityTrait, QueryTrait};
+
+fn main() {
+    assert_eq!(
+        Container::find().find_with_related(Content).build(DbBackend::MySql).to_string(),
+        [
+            "SELECT `container`.`db_id` AS `A_db_id`, `content`.`id` AS `B_id`, `content`.`container_id` AS `B_container_id`",
+            "FROM `container`",
+            "LEFT JOIN `content` ON `container`.`db_id` = `content`.`container_id`",
+            "ORDER BY `container`.`db_id` ASC",
+        ].join(" ")
+    );
+}

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 {

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();
 

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);