diff --git a/sea-orm-macros/src/derives/from_query_result.rs b/sea-orm-macros/src/derives/from_query_result.rs index de511f4ed..9795eb692 100644 --- a/sea-orm-macros/src/derives/from_query_result.rs +++ b/sea-orm-macros/src/derives/from_query_result.rs @@ -1,20 +1,17 @@ +use std::collections::{HashMap, hash_map::Entry}; + use super::util::GetMeta; use proc_macro2::{Ident, TokenStream}; -use quote::{ToTokens, format_ident, quote, quote_spanned}; +use quote::{ToTokens, quote}; use syn::{ - Data, DataStruct, DeriveInput, Fields, Generics, Meta, ext::IdentExt, punctuated::Punctuated, - token::Comma, + Data, DataStruct, DeriveInput, Error, Fields, Generics, Meta, ext::IdentExt, + punctuated::Punctuated, token::Comma, }; -#[derive(Debug)] -enum Error { - InputNotStruct, -} - pub(super) enum ItemType { Flat, Skip, - Nested, + Nested { prefix: Option }, } pub(super) struct DeriveFromQueryResult { @@ -63,13 +60,16 @@ impl ToTokens for TryFromQueryResultCheck<'_> { let #ident = std::default::Default::default(); }); } - ItemType::Nested => { - let prefix = if self.0 { - let name = ident.unraw().to_string(); - quote! { &format!("{pre}{}_", #name) } - } else { - quote! { pre } + ItemType::Nested { prefix } => { + let prefix = match (self.0, prefix) { + (_, Some(p)) => quote! { &format!("{pre}{}", #p) }, + (true, None) => { + let name = ident.unraw().to_string(); + quote! { &format!("{pre}{}_", #name) } + } + (false, None) => quote! { pre }, }; + tokens.extend(quote! { let #ident = match sea_orm::FromQueryResult::from_query_result_nullable(row, #prefix) { Err(v @ sea_orm::TryGetError::DbErr(_)) => { @@ -90,7 +90,7 @@ impl ToTokens for TryFromQueryResultAssignment<'_> { let FromQueryResultItem { ident, typ, .. } = self.0; match typ { - ItemType::Flat | ItemType::Nested => { + ItemType::Flat | ItemType::Nested { .. } => { tokens.extend(quote! { #ident: #ident?, }); @@ -118,10 +118,17 @@ impl DeriveFromQueryResult { fields: Fields::Named(named), .. }) => named.named, - _ => return Err(Error::InputNotStruct), + _ => { + return Err(Error::new( + ident.span(), + "you can only derive `FromQueryResult` on named struct", + )); + } }; let mut fields = Vec::with_capacity(parsed_fields.len()); + let mut seen_nested: HashMap<(syn::Type, Option), TokenStream> = HashMap::new(); + for parsed_field in parsed_fields { let mut typ = ItemType::Flat; let mut alias = None; @@ -135,16 +142,59 @@ impl DeriveFromQueryResult { if meta.exists("skip") { typ = ItemType::Skip; } else if meta.exists("nested") { - typ = ItemType::Nested; - } else if let Some(alias_) = meta.get_as_kv("from_alias") { - alias = Some(alias_); + typ = ItemType::Nested { prefix: None }; + } else if let Some(list) = meta.get_list_args("nested") { + let mut prefix = None; + + for m in list.iter() { + match m.get_as_kv("prefix") { + Some(p) => prefix = (!p.is_empty()).then_some(p), + None => { + return Err(Error::new_spanned( + m, + "invalid nested attribute, expected `prefix = \"...\"`", + )); + } + } + } + + typ = ItemType::Nested { prefix }; } else { - alias = meta.get_as_kv("alias"); + alias = meta + .get_as_kv("from_alias") + .or_else(|| meta.get_as_kv("alias")); } } } } - let ident = format_ident!("{}", parsed_field.ident.unwrap().to_string()); + + let field_tokens = parsed_field.to_token_stream(); + let ident = parsed_field.ident.unwrap(); + + if let ItemType::Nested { ref prefix } = typ { + let key = (parsed_field.ty, prefix.clone()); + match seen_nested.entry(key) { + Entry::Occupied(e) => { + let msg = match prefix { + Some(p) => format!( + "multiple nested fields with the same type share prefix \"{p}\"" + ), + None => { + "multiple nested fields with the same type must have a `prefix`: \ + use `#[sea_orm(nested(prefix = \"...\"))]`" + .to_string() + } + }; + let mut err = Error::new_spanned(&field_tokens, msg); + err.combine(Error::new_spanned(e.get(), "first defined here")); + return Err(err); + } + Entry::Vacant(e) => { + e.insert(field_tokens); + } + } + } + fields.push(FromQueryResultItem { typ, ident, alias }); } @@ -194,12 +244,5 @@ impl DeriveFromQueryResult { } pub fn expand_derive_from_query_result(input: DeriveInput) -> syn::Result { - let ident_span = input.ident.span(); - - match DeriveFromQueryResult::new(input) { - Ok(partial_model) => partial_model.expand(), - Err(Error::InputNotStruct) => Ok(quote_spanned! { - ident_span => compile_error!("you can only derive `FromQueryResult` on named struct"); - }), - } + DeriveFromQueryResult::new(input)?.expand() } diff --git a/sea-orm-macros/src/derives/partial_model.rs b/sea-orm-macros/src/derives/partial_model.rs index a16dbfc9a..455efa4a2 100644 --- a/sea-orm-macros/src/derives/partial_model.rs +++ b/sea-orm-macros/src/derives/partial_model.rs @@ -1,6 +1,8 @@ +use std::collections::{HashMap, hash_map::Entry}; + use heck::ToUpperCamelCase; use proc_macro2::{Span, TokenStream}; -use quote::{format_ident, quote, quote_spanned}; +use quote::{ToTokens, format_ident, quote, quote_spanned}; use syn::{ Expr, Meta, Type, ext::IdentExt, punctuated::Punctuated, spanned::Spanned, token::Comma, }; @@ -37,6 +39,7 @@ enum ColumnAs { typ: Type, field: syn::Ident, alias: Option, + prefix: Option, }, Skip(syn::Ident), } @@ -105,6 +108,7 @@ impl DerivePartialModel { } let mut column_as_list = Vec::with_capacity(fields.len()); + let mut seen_nested: HashMap<(syn::Type, Option), TokenStream> = HashMap::new(); for field in fields { let field_span = field.span(); @@ -113,6 +117,7 @@ impl DerivePartialModel { let mut from_expr = None; let mut nested = false; let mut nested_alias = None; + let mut nested_prefix = None; let mut skip = false; for attr in field.attrs.iter() { @@ -127,6 +132,19 @@ impl DerivePartialModel { skip = true; } else if meta.exists("nested") { nested = true; + } else if let Some(list) = meta.get_list_args("nested") { + nested = true; + for m in list.iter() { + match m.get_as_kv("prefix") { + Some(p) => nested_prefix = Some(p), + None => { + return Err(Error::Syn(syn::Error::new_spanned( + m, + "invalid nested attribute, expected `prefix = \"...\"`", + ))); + } + } + } } else if let Some(s) = meta.get_as_kv("from_col") { from_col = Some(format_ident!("{}", s.to_upper_camel_case())); } else if let Some(s) = meta.get_as_kv("from_expr") { @@ -138,6 +156,7 @@ impl DerivePartialModel { } } + let field_tokens = field.to_token_stream(); let field_name = field.ident.unwrap(); let col_as = match (from_col, from_expr, nested) { @@ -155,11 +174,36 @@ impl DerivePartialModel { expr, field: field_name, }, - (None, None, true) => ColumnAs::Nested { - typ: field.ty, - field: field_name, - alias: nested_alias, - }, + (None, None, true) => { + let key = (field.ty.clone(), nested_prefix.clone()); + match seen_nested.entry(key) { + Entry::Occupied(e) => { + let msg = match nested_prefix { + Some(p) => format!( + "multiple nested fields with the same type share prefix \"{p}\"" + ), + None => { + "multiple nested fields with the same type must have a `prefix`: \ + use `#[sea_orm(nested(prefix = \"...\"))]`" + .to_string() + } + }; + let mut err = syn::Error::new_spanned(&field_tokens, msg); + err.combine(syn::Error::new_spanned(e.get(), "first defined here")); + return Err(Error::Syn(err)); + } + Entry::Vacant(e) => { + e.insert(field_tokens); + } + } + + ColumnAs::Nested { + typ: field.ty, + field: field_name, + alias: nested_alias, + prefix: nested_prefix, + } + } (None, None, false) => { if entity.is_none() { return Err(Error::EntityNotSpecified); @@ -201,7 +245,9 @@ impl DerivePartialModel { .iter() .map(|col_as| FromQueryResultItem { typ: match col_as { - ColumnAs::Nested { .. } => FqrItemType::Nested, + ColumnAs::Nested { prefix, .. } => FqrItemType::Nested { + prefix: prefix.clone(), + }, ColumnAs::Skip(_) => FqrItemType::Skip, _ => FqrItemType::Flat, }, @@ -320,22 +366,37 @@ impl DerivePartialModel { }; ) } - ColumnAs::Nested { typ, field, alias } => { - let field = field.unraw().to_string(); + ColumnAs::Nested { + typ, + field, + alias, + prefix, + } => { + let field_str = field.unraw().to_string(); let alias_ref: Option<&str> = alias.as_deref(); let alias_arg = match alias_ref { Some(s) => quote! { Some(#s) }, None => quote! { None }, }; - quote!(let #select_ident = - <#typ as sea_orm::PartialModelTrait>::select_cols_nested(#select_ident, + let prefix_expr = match prefix { + Some(p) => quote! { Some(&if let Some(prefix) = pre { - format!("{prefix}{}_", #field) - } else { - format!("{}_", #field) - } - ), - #alias_arg + format!("{prefix}{}", #p) + } else { + #p.to_string() + }) + }, + None => quote! { + Some(&if let Some(prefix) = pre { + format!("{prefix}{}_", #field_str) + } else { + format!("{}_", #field_str) + }) + }, + }; + quote!(let #select_ident = + <#typ as sea_orm::PartialModelTrait>::select_cols_nested( + #select_ident, #prefix_expr, #alias_arg ); ) } diff --git a/sea-orm-macros/src/derives/util.rs b/sea-orm-macros/src/derives/util.rs index 276435e67..1db7a0424 100644 --- a/sea-orm-macros/src/derives/util.rs +++ b/sea-orm-macros/src/derives/util.rs @@ -199,6 +199,7 @@ pub(crate) trait GetMeta { fn exists(&self, k: &str) -> bool; fn get_as_kv(&self, k: &str) -> Option; fn get_as_kv_with_ident(&self) -> Option<(Ident, String)>; + fn get_list_args(&self, name: &str) -> Option>; } impl GetMeta for Meta { @@ -247,6 +248,15 @@ impl GetMeta for Meta { path.get_ident() .map(|ident| (ident.clone(), litstr.value())) } + + fn get_list_args(&self, name: &str) -> Option> { + match self { + Meta::List(list) if list.path.is_ident(name) => list + .parse_args_with(Punctuated::::parse_terminated) + .ok(), + _ => None, + } + } } #[cfg(test)] diff --git a/sea-orm-macros/src/lib.rs b/sea-orm-macros/src/lib.rs index 0abe34adb..3b5cc9e9c 100644 --- a/sea-orm-macros/src/lib.rs +++ b/sea-orm-macros/src/lib.rs @@ -708,7 +708,7 @@ pub fn derive_active_enum(input: TokenStream) -> TokenStream { /// ### Attributes /// /// - `skip`: will not try to pull this field from the query result. And set it to the default value of the type. -/// - `nested`: allows nesting models. can be any type that implements `FromQueryResult` +/// - `nested`: allows nesting models. can be any type that implements `FromQueryResult`. supports `nested(prefix = "...")` to set an explicit column prefix. /// - `alias` / `from_alias`: get the value from this column alias /// /// ### Usage @@ -747,6 +747,9 @@ pub fn derive_active_enum(input: TokenStream) -> TokenStream { /// price: Decimal, /// #[sea_orm(nested)] /// baker: Option, +/// // prefix is optional, useful when multiple fields share the same type +/// #[sea_orm(nested(prefix = "reviewer_"))] +/// reviewer: Option, /// } /// ``` #[cfg(feature = "derive")]