diff --git a/sea-orm-macros/src/derives/from_query_result.rs b/sea-orm-macros/src/derives/from_query_result.rs index 999db62fe..a072827b2 100644 --- a/sea-orm-macros/src/derives/from_query_result.rs +++ b/sea-orm-macros/src/derives/from_query_result.rs @@ -1,6 +1,29 @@ +use self::util::GetMeta; use proc_macro2::{Ident, TokenStream}; -use quote::{format_ident, quote, quote_spanned}; -use syn::{ext::IdentExt, Data, DataStruct, Field, Fields, Generics}; +use quote::{format_ident, quote, quote_spanned, ToTokens}; +use syn::{ + ext::IdentExt, punctuated::Punctuated, token::Comma, Data, DataStruct, Fields, Generics, Meta, +}; + +pub struct FromQueryResultItem { + pub skip: bool, + pub ident: Ident, +} +impl ToTokens for FromQueryResultItem { + fn to_tokens(&self, tokens: &mut TokenStream) { + let Self { ident, skip } = self; + if *skip { + tokens.extend(quote! { + #ident: std::default::Default::default(), + }); + } else { + let name = ident.unraw().to_string(); + tokens.extend(quote! { + #ident: row.try_get(pre, #name)?, + }); + } + } +} /// Method to derive a [QueryResult](sea_orm::QueryResult) pub fn expand_derive_from_query_result( @@ -19,20 +42,23 @@ pub fn expand_derive_from_query_result( }) } }; + let mut field = Vec::with_capacity(fields.len()); - let field: Vec = fields - .into_iter() - .map(|Field { ident, .. }| format_ident!("{}", ident.unwrap().to_string())) - .collect(); - - let name: Vec = field - .iter() - .map(|f| { - let s = f.unraw().to_string(); - quote! { #s } - }) - .collect(); - + for parsed_field in fields.into_iter() { + let mut skip = false; + for attr in parsed_field.attrs.iter() { + if !attr.path().is_ident("sea_orm") { + continue; + } + if let Ok(list) = attr.parse_args_with(Punctuated::::parse_terminated) { + for meta in list.iter() { + skip = meta.exists("skip"); + } + } + } + let ident = format_ident!("{}", parsed_field.ident.unwrap().to_string()); + field.push(FromQueryResultItem { skip, ident }); + } let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); Ok(quote!( @@ -40,9 +66,25 @@ pub fn expand_derive_from_query_result( impl #impl_generics sea_orm::FromQueryResult for #ident #ty_generics #where_clause { fn from_query_result(row: &sea_orm::QueryResult, pre: &str) -> std::result::Result { Ok(Self { - #(#field: row.try_get(pre, #name)?),* + #(#field)* }) } } )) } +mod util { + use syn::Meta; + + pub(super) trait GetMeta { + fn exists(&self, k: &str) -> bool; + } + + impl GetMeta for Meta { + fn exists(&self, k: &str) -> bool { + let Meta::Path(path) = self else { + return false; + }; + path.is_ident(k) + } + } +} diff --git a/sea-orm-macros/src/lib.rs b/sea-orm-macros/src/lib.rs index 85a9fee62..52f2e734a 100644 --- a/sea-orm-macros/src/lib.rs +++ b/sea-orm-macros/src/lib.rs @@ -579,6 +579,9 @@ pub fn derive_active_enum(input: TokenStream) -> TokenStream { /// Convert a query result into the corresponding Model. /// +/// ### Attributes +/// - `skip`: Will not try to pull this field from the query result. And set it to the default value of the type. +/// /// ### Usage /// /// ``` @@ -588,10 +591,12 @@ pub fn derive_active_enum(input: TokenStream) -> TokenStream { /// struct SelectResult { /// name: String, /// num_of_fruits: i32, +/// #[sea_orm(skip)] +/// skip_me: i32, /// } /// ``` #[cfg(feature = "derive")] -#[proc_macro_derive(FromQueryResult)] +#[proc_macro_derive(FromQueryResult, attributes(sea_orm))] pub fn derive_from_query_result(input: TokenStream) -> TokenStream { let DeriveInput { ident, diff --git a/tests/derive_tests.rs b/tests/derive_tests.rs index 20e49d06b..bd99a3b46 100644 --- a/tests/derive_tests.rs +++ b/tests/derive_tests.rs @@ -56,3 +56,10 @@ where { _foo: T::Item, } + +#[derive(FromQueryResult)] +struct FromQueryAttributeTests { + #[sea_orm(skip)] + _foo: i32, + _bar: String, +}