caisin/sea-orm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update DerivePrimaryKey

Browse Source
This commit is contained in:
Billy Chan 2021-09-27 18:02:12 +08:00
parent f79a417150
commit 4f090d192b
No known key found for this signature in database
GPG Key ID: A2D690CAC7DF3CC7
1 changed files with 28 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
sea-orm-macros/src/derives/primary_key.rs
View File

@ -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
}
}
))
}