Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 73 additions & 30 deletions sea-orm-macros/src/derives/from_query_result.rs
Original file line number Diff line number Diff line change
@@ -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<String> },
}

pub(super) struct DeriveFromQueryResult {
Expand Down Expand Up @@ -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(_)) => {
Expand All @@ -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?,
});
Expand Down Expand Up @@ -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<String>), TokenStream> = HashMap::new();

for parsed_field in parsed_fields {
let mut typ = ItemType::Flat;
let mut alias = None;
Expand All @@ -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 });
}

Expand Down Expand Up @@ -194,12 +244,5 @@ impl DeriveFromQueryResult {
}

pub fn expand_derive_from_query_result(input: DeriveInput) -> syn::Result<TokenStream> {
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()
}
95 changes: 78 additions & 17 deletions sea-orm-macros/src/derives/partial_model.rs
Original file line number Diff line number Diff line change
@@ -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,
};
Expand Down Expand Up @@ -37,6 +39,7 @@ enum ColumnAs {
typ: Type,
field: syn::Ident,
alias: Option<String>,
prefix: Option<String>,
},
Skip(syn::Ident),
}
Expand Down Expand Up @@ -105,6 +108,7 @@ impl DerivePartialModel {
}

let mut column_as_list = Vec::with_capacity(fields.len());
let mut seen_nested: HashMap<(syn::Type, Option<String>), TokenStream> = HashMap::new();

for field in fields {
let field_span = field.span();
Expand All @@ -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() {
Expand All @@ -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") {
Expand All @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -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,
},
Expand Down Expand Up @@ -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
);
)
}
Expand Down
10 changes: 10 additions & 0 deletions sea-orm-macros/src/derives/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,7 @@ pub(crate) trait GetMeta {
fn exists(&self, k: &str) -> bool;
fn get_as_kv(&self, k: &str) -> Option<String>;
fn get_as_kv_with_ident(&self) -> Option<(Ident, String)>;
fn get_list_args(&self, name: &str) -> Option<Punctuated<Meta, Comma>>;
}

impl GetMeta for Meta {
Expand Down Expand Up @@ -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<Punctuated<Meta, Comma>> {
match self {
Meta::List(list) if list.path.is_ident(name) => list
.parse_args_with(Punctuated::<Meta, Comma>::parse_terminated)
.ok(),
_ => None,
}
}
}

#[cfg(test)]
Expand Down
5 changes: 4 additions & 1 deletion sea-orm-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -747,6 +747,9 @@ pub fn derive_active_enum(input: TokenStream) -> TokenStream {
/// price: Decimal,
/// #[sea_orm(nested)]
/// baker: Option<cakes_bakers::Model>,
/// // prefix is optional, useful when multiple fields share the same type
/// #[sea_orm(nested(prefix = "reviewer_"))]
/// reviewer: Option<cakes_bakers::Model>,
/// }
/// ```
#[cfg(feature = "derive")]
Expand Down