From 07bb07f753dc7e121a15001bfea4ccd8ba2f6164 Mon Sep 17 00:00:00 2001 From: Wyatt Herkamp Date: Thu, 2 Nov 2023 09:13:38 -0400 Subject: [PATCH 1/3] Add #[sea_orm(skip)] for FromQueryResult macro --- .../src/derives/from_query_result.rs | 74 +++++++++++++++---- sea-orm-macros/src/lib.rs | 7 +- tests/derive_tests.rs | 7 ++ 3 files changed, 71 insertions(+), 17 deletions(-) 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, +} From b7226bb44dbdbab9782524feb0303c4c166d3050 Mon Sep 17 00:00:00 2001 From: Wyatt Herkamp Date: Fri, 3 Nov 2023 10:04:54 -0400 Subject: [PATCH 2/3] Update sea-orm-rocket --- sea-orm-rocket/codegen/src/database.rs | 2 +- sea-orm-rocket/lib/src/database.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sea-orm-rocket/codegen/src/database.rs b/sea-orm-rocket/codegen/src/database.rs index 360c84742..c1aba5f3f 100644 --- a/sea-orm-rocket/codegen/src/database.rs +++ b/sea-orm-rocket/codegen/src/database.rs @@ -64,7 +64,7 @@ pub fn derive_database(input: TokenStream) -> TokenStream { ) -> rocket::request::Outcome { match #db_ty::fetch(req.rocket()) { Some(db) => rocket::outcome::Outcome::Success(db), - None => rocket::outcome::Outcome::Failure(( + None => rocket::outcome::Outcome::Error(( rocket::http::Status::InternalServerError, ())) } } diff --git a/sea-orm-rocket/lib/src/database.rs b/sea-orm-rocket/lib/src/database.rs index 54f6da34d..fe275e185 100644 --- a/sea-orm-rocket/lib/src/database.rs +++ b/sea-orm-rocket/lib/src/database.rs @@ -131,10 +131,10 @@ pub trait Database: } let dbtype = std::any::type_name::(); - let fairing = Paint::default(format!("{dbtype}::init()")).bold(); + let fairing = Paint::new(format!("{dbtype}::init()")).bold(); error!( "Attempted to fetch unattached database `{}`.", - Paint::default(dbtype).bold() + Paint::new(dbtype).bold() ); info_!( "`{}` fairing must be attached prior to using this database.", @@ -261,7 +261,7 @@ impl<'r, D: Database> FromRequest<'r> for Connection<'r, D> { async fn from_request(req: &'r Request<'_>) -> Outcome { match D::fetch(req.rocket()) { Some(pool) => Outcome::Success(Connection(pool.borrow())), - None => Outcome::Failure((Status::InternalServerError, None)), + None => Outcome::Error((Status::InternalServerError, None)), } } } From cf93479d62b9c9c0476d8c26c8d73e5518ea80de Mon Sep 17 00:00:00 2001 From: Wyatt Herkamp Date: Fri, 3 Nov 2023 10:44:39 -0400 Subject: [PATCH 3/3] Revert "Update sea-orm-rocket" This reverts commit b7226bb44dbdbab9782524feb0303c4c166d3050. --- sea-orm-rocket/codegen/src/database.rs | 2 +- sea-orm-rocket/lib/src/database.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sea-orm-rocket/codegen/src/database.rs b/sea-orm-rocket/codegen/src/database.rs index c1aba5f3f..360c84742 100644 --- a/sea-orm-rocket/codegen/src/database.rs +++ b/sea-orm-rocket/codegen/src/database.rs @@ -64,7 +64,7 @@ pub fn derive_database(input: TokenStream) -> TokenStream { ) -> rocket::request::Outcome { match #db_ty::fetch(req.rocket()) { Some(db) => rocket::outcome::Outcome::Success(db), - None => rocket::outcome::Outcome::Error(( + None => rocket::outcome::Outcome::Failure(( rocket::http::Status::InternalServerError, ())) } } diff --git a/sea-orm-rocket/lib/src/database.rs b/sea-orm-rocket/lib/src/database.rs index fe275e185..54f6da34d 100644 --- a/sea-orm-rocket/lib/src/database.rs +++ b/sea-orm-rocket/lib/src/database.rs @@ -131,10 +131,10 @@ pub trait Database: } let dbtype = std::any::type_name::(); - let fairing = Paint::new(format!("{dbtype}::init()")).bold(); + let fairing = Paint::default(format!("{dbtype}::init()")).bold(); error!( "Attempted to fetch unattached database `{}`.", - Paint::new(dbtype).bold() + Paint::default(dbtype).bold() ); info_!( "`{}` fairing must be attached prior to using this database.", @@ -261,7 +261,7 @@ impl<'r, D: Database> FromRequest<'r> for Connection<'r, D> { async fn from_request(req: &'r Request<'_>) -> Outcome { match D::fetch(req.rocket()) { Some(pool) => Outcome::Success(Connection(pool.borrow())), - None => Outcome::Error((Status::InternalServerError, None)), + None => Outcome::Failure((Status::InternalServerError, None)), } } }