sea_orm_macros/derives/
from_query_result.rs1use 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
32struct 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}