sea_orm_macros/derives/
from_query_result.rs

1use super::util::GetMeta;
2use proc_macro2::{Ident, TokenStream};
3use quote::{format_ident, quote, quote_spanned, ToTokens};
4use syn::{
5    ext::IdentExt, punctuated::Punctuated, token::Comma, Data, DataStruct, DeriveInput, Fields,
6    Generics, Meta,
7};
8
9#[derive(Debug)]
10enum Error {
11    InputNotStruct,
12}
13
14pub(super) enum ItemType {
15    Flat,
16    Skip,
17    Nested,
18}
19
20pub(super) struct DeriveFromQueryResult {
21    pub ident: syn::Ident,
22    pub generics: Generics,
23    pub fields: Vec<FromQueryResultItem>,
24}
25
26pub(super) struct FromQueryResultItem {
27    pub typ: ItemType,
28    pub ident: Ident,
29    pub alias: Option<String>,
30}
31
32/// Initially, we try to obtain the value for each field and check if it is an ordinary DB error
33/// (which we return immediatly), or a null error.
34///
35/// ### Background
36///
37/// Null errors do not necessarily mean that the deserialization as a whole fails,
38/// since structs embedding the current one might have wrapped the current one in an `Option`.
39/// In this case, we do not want to swallow other errors, which are very likely to actually be
40/// programming errors that should be noticed (and fixed).
41struct TryFromQueryResultCheck<'a>(bool, &'a FromQueryResultItem);
42
43impl ToTokens for TryFromQueryResultCheck<'_> {
44    fn to_tokens(&self, tokens: &mut TokenStream) {
45        let FromQueryResultItem { ident, typ, alias } = self.1;
46
47        match typ {
48            ItemType::Flat => {
49                let name = alias
50                    .to_owned()
51                    .unwrap_or_else(|| ident.unraw().to_string());
52                tokens.extend(quote! {
53                    let #ident = match row.try_get_nullable(pre, #name) {
54                        Err(v @ sea_orm::TryGetError::DbErr(_)) => {
55                            return Err(v);
56                        }
57                        v => v,
58                    };
59                });
60            }
61            ItemType::Skip => {
62                tokens.extend(quote! {
63                    let #ident = std::default::Default::default();
64                });
65            }
66            ItemType::Nested => {
67                let prefix = if self.0 {
68                    let name = ident.unraw().to_string();
69                    quote! { &format!("{pre}{}_", #name) }
70                } else {
71                    quote! { pre }
72                };
73                tokens.extend(quote! {
74                    let #ident = match sea_orm::FromQueryResult::from_query_result_nullable(row, #prefix) {
75                        Err(v @ sea_orm::TryGetError::DbErr(_)) => {
76                            return Err(v);
77                        }
78                        v => v,
79                    };
80                });
81            }
82        }
83    }
84}
85
86struct TryFromQueryResultAssignment<'a>(&'a FromQueryResultItem);
87
88impl ToTokens for TryFromQueryResultAssignment<'_> {
89    fn to_tokens(&self, tokens: &mut TokenStream) {
90        let FromQueryResultItem { ident, typ, .. } = self.0;
91
92        match typ {
93            ItemType::Flat | ItemType::Nested => {
94                tokens.extend(quote! {
95                    #ident: #ident?,
96                });
97            }
98            ItemType::Skip => {
99                tokens.extend(quote! {
100                    #ident,
101                });
102            }
103        }
104    }
105}
106
107impl DeriveFromQueryResult {
108    fn new(
109        DeriveInput {
110            ident,
111            data,
112            generics,
113            ..
114        }: DeriveInput,
115    ) -> Result<Self, Error> {
116        let parsed_fields = match data {
117            Data::Struct(DataStruct {
118                fields: Fields::Named(named),
119                ..
120            }) => named.named,
121            _ => return Err(Error::InputNotStruct),
122        };
123
124        let mut fields = Vec::with_capacity(parsed_fields.len());
125        for parsed_field in parsed_fields {
126            let mut typ = ItemType::Flat;
127            let mut alias = None;
128            for attr in parsed_field.attrs.iter() {
129                if !attr.path().is_ident("sea_orm") {
130                    continue;
131                }
132                if let Ok(list) = attr.parse_args_with(Punctuated::<Meta, Comma>::parse_terminated)
133                {
134                    for meta in list.iter() {
135                        if meta.exists("skip") {
136                            typ = ItemType::Skip;
137                        } else if meta.exists("nested") {
138                            typ = ItemType::Nested;
139                        } else {
140                            alias = meta.get_as_kv("from_alias");
141                        }
142                    }
143                }
144            }
145            let ident = format_ident!("{}", parsed_field.ident.unwrap().to_string());
146            fields.push(FromQueryResultItem { typ, ident, alias });
147        }
148
149        Ok(Self {
150            ident,
151            generics,
152            fields,
153        })
154    }
155
156    fn expand(&self) -> syn::Result<TokenStream> {
157        Ok(self.impl_from_query_result(false))
158    }
159
160    pub(super) fn impl_from_query_result(&self, prefix: bool) -> TokenStream {
161        let Self {
162            ident,
163            generics,
164            fields,
165        } = self;
166
167        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
168
169        let ident_try_init: Vec<_> = fields
170            .iter()
171            .map(|s| TryFromQueryResultCheck(prefix, s))
172            .collect();
173        let ident_try_assign: Vec<_> = fields.iter().map(TryFromQueryResultAssignment).collect();
174
175        quote!(
176            #[automatically_derived]
177            impl #impl_generics sea_orm::FromQueryResult for #ident #ty_generics #where_clause {
178                fn from_query_result(row: &sea_orm::QueryResult, pre: &str) -> std::result::Result<Self, sea_orm::DbErr> {
179                    Ok(Self::from_query_result_nullable(row, pre)?)
180                }
181
182                fn from_query_result_nullable(row: &sea_orm::QueryResult, pre: &str) -> std::result::Result<Self, sea_orm::TryGetError> {
183                    #(#ident_try_init)*
184
185                    Ok(Self {
186                        #(#ident_try_assign)*
187                    })
188                }
189            }
190        )
191    }
192}
193
194pub fn expand_derive_from_query_result(input: DeriveInput) -> syn::Result<TokenStream> {
195    let ident_span = input.ident.span();
196
197    match DeriveFromQueryResult::new(input) {
198        Ok(partial_model) => partial_model.expand(),
199        Err(Error::InputNotStruct) => Ok(quote_spanned! {
200            ident_span => compile_error!("you can only derive `FromQueryResult` on named struct");
201        }),
202    }
203}