1use proc_macro::TokenStream;
29use proc_macro2::TokenStream as TokenStream2;
30use quote::quote;
31use syn::{parse_macro_input, Data, DeriveInput, Fields, LitStr};
32
33#[proc_macro_derive(Entity, attributes(orm))]
34pub fn derive_entity(input: TokenStream) -> TokenStream {
35 let input = parse_macro_input!(input as DeriveInput);
36 let ident = input.ident.clone();
37
38 struct IndexDef {
39 columns: Vec<String>,
40 unique: bool,
41 }
42
43 let mut table = to_snake_case(&ident.to_string());
45 let mut composite: Vec<IndexDef> = Vec::new();
46 for attr in &input.attrs {
47 if attr.path().is_ident("orm") {
48 let _ = attr.parse_nested_meta(|meta| {
49 if meta.path.is_ident("table") {
50 table = meta.value()?.parse::<LitStr>()?.value();
51 } else if meta.path.is_ident("index") {
52 composite.push(IndexDef {
53 columns: split_columns(&meta.value()?.parse::<LitStr>()?.value()),
54 unique: false,
55 });
56 } else if meta.path.is_ident("unique") {
57 composite.push(IndexDef {
58 columns: split_columns(&meta.value()?.parse::<LitStr>()?.value()),
59 unique: true,
60 });
61 }
62 Ok(())
63 });
64 }
65 }
66
67 let fields = match &input.data {
68 Data::Struct(s) => match &s.fields {
69 Fields::Named(n) => n.named.clone(),
70 _ => return compile_err("Entity requires named struct fields"),
71 },
72 _ => return compile_err("Entity can only be derived for a struct"),
73 };
74
75 struct FieldMeta {
76 ident: syn::Ident,
77 ty: syn::Type,
78 column: String,
79 pk: bool,
80 auto_increment: bool,
81 unique: bool,
82 index: bool,
83 shard: bool,
84 references: Option<(String, String)>, on_delete: Option<String>,
86 on_update: Option<String>,
87 }
88
89 let mut metas: Vec<FieldMeta> = Vec::new();
90 let mut pk: Option<(syn::Ident, String)> = None;
91 for f in &fields {
92 let fident = f.ident.clone().unwrap();
93 let mut column = fident.to_string();
94 let mut is_pk = false;
95 let mut auto_increment = false;
96 let mut unique = false;
97 let mut index = false;
98 let mut shard = false;
99 let mut references = None;
100 let mut on_delete = None;
101 let mut on_update = None;
102 for attr in &f.attrs {
103 if attr.path().is_ident("orm") {
104 let _ = attr.parse_nested_meta(|meta| {
105 if meta.path.is_ident("pk") {
106 is_pk = true;
107 } else if meta.path.is_ident("auto_increment") {
108 auto_increment = true;
109 } else if meta.path.is_ident("unique") {
110 unique = true;
111 } else if meta.path.is_ident("index") {
112 index = true;
113 } else if meta.path.is_ident("shard_key") {
114 shard = true;
115 } else if meta.path.is_ident("column") {
116 column = meta.value()?.parse::<LitStr>()?.value();
117 } else if meta.path.is_ident("references") {
118 let v = meta.value()?.parse::<LitStr>()?.value();
119 references = Some(parse_references(&v));
120 } else if meta.path.is_ident("on_delete") {
121 on_delete = Some(meta.value()?.parse::<LitStr>()?.value());
122 } else if meta.path.is_ident("on_update") {
123 on_update = Some(meta.value()?.parse::<LitStr>()?.value());
124 }
125 Ok(())
126 });
127 }
128 }
129 if is_pk {
130 pk = Some((fident.clone(), column.clone()));
131 }
132 metas.push(FieldMeta {
133 ident: fident,
134 ty: f.ty.clone(),
135 column,
136 pk: is_pk,
137 auto_increment,
138 unique,
139 index,
140 shard,
141 references,
142 on_delete,
143 on_update,
144 });
145 }
146
147 let Some((pk_ident, pk_col)) = pk else {
148 return compile_err("Entity needs exactly one field marked `#[orm(pk)]`");
149 };
150
151 let col_specs = metas.iter().map(|m| {
152 let name = &m.column;
153 let ty = &m.ty;
154 let is_pk = m.pk;
155 let ai = m.auto_increment;
156 let uniq = m.unique;
157 quote! {
158 ::rainier_orm::ColumnSpec {
159 name: #name,
160 ty: <#ty as ::rainier_orm::SqlType>::COLUMN_TYPE,
161 nullable: <#ty as ::rainier_orm::SqlType>::NULLABLE,
162 pk: #is_pk,
163 auto_increment: #ai,
164 unique: #uniq,
165 }
166 }
167 });
168
169 let mut index_specs: Vec<TokenStream2> = Vec::new();
171 for m in &metas {
172 if m.index {
173 let name = format!("idx_{table}_{}", m.column);
174 let col = &m.column;
175 index_specs.push(quote! {
176 ::rainier_orm::IndexSpec { name: #name, columns: &[#col], unique: false }
177 });
178 }
179 }
180 for ix in &composite {
181 let joined = ix.columns.join("_");
182 let prefix = if ix.unique { "uq" } else { "idx" };
183 let name = format!("{prefix}_{table}_{joined}");
184 let cols = &ix.columns;
185 let uniq = ix.unique;
186 index_specs.push(quote! {
187 ::rainier_orm::IndexSpec { name: #name, columns: &[#(#cols),*], unique: #uniq }
188 });
189 }
190
191 let fk_specs = metas.iter().filter_map(|m| {
193 let (ftable, fcol) = m.references.as_ref()?;
194 let name = format!("fk_{table}_{}", m.column);
195 let col = &m.column;
196 let on_delete = ref_action(m.on_delete.as_deref());
197 let on_update = ref_action(m.on_update.as_deref());
198 Some(quote! {
199 ::rainier_orm::ForeignKeySpec {
200 name: #name,
201 columns: &[#col],
202 foreign_table: #ftable,
203 foreign_columns: &[#fcol],
204 on_delete: #on_delete,
205 on_update: #on_update,
206 }
207 })
208 });
209
210 let shard_cols = metas.iter().filter(|m| m.shard).map(|m| {
212 let c = &m.column;
213 quote! { #c }
214 });
215
216 let from_row_fields = metas.iter().map(|m| {
217 let fident = &m.ident;
218 let ty = &m.ty;
219 let col = &m.column;
220 quote! { #fident: <#ty as ::rainier_orm::FromColumn>::from_column(row, #col)? }
221 });
222
223 let insert_vals = metas.iter().filter(|m| !m.auto_increment).map(|m| {
225 let fident = &m.ident;
226 let col = &m.column;
227 quote! { (#col, ::rainier_orm::ToColumn::to_value(&self.#fident)) }
228 });
229 let update_vals = metas.iter().filter(|m| !m.pk).map(|m| {
231 let fident = &m.ident;
232 let col = &m.column;
233 quote! { (#col, ::rainier_orm::ToColumn::to_value(&self.#fident)) }
234 });
235
236 quote! {
237 impl ::rainier_orm::Entity for #ident {
238 fn table() -> &'static str { #table }
239 fn columns() -> &'static [::rainier_orm::ColumnSpec] {
240 &[ #(#col_specs),* ]
241 }
242 fn primary_key() -> &'static str { #pk_col }
243 fn indexes() -> &'static [::rainier_orm::IndexSpec] {
244 &[ #(#index_specs),* ]
245 }
246 fn foreign_keys() -> &'static [::rainier_orm::ForeignKeySpec] {
247 &[ #(#fk_specs),* ]
248 }
249 fn shard_columns() -> &'static [&'static str] {
250 &[ #(#shard_cols),* ]
251 }
252 fn from_row(row: &dyn ::rainier_orm::Row)
253 -> ::core::result::Result<Self, ::rainier_orm::Error>
254 {
255 ::core::result::Result::Ok(Self { #(#from_row_fields),* })
256 }
257 fn insert_values(&self)
258 -> ::std::vec::Vec<(&'static str, ::rainier_orm::sea_query::Value)>
259 {
260 ::std::vec![ #(#insert_vals),* ]
261 }
262 fn update_values(&self)
263 -> ::std::vec::Vec<(&'static str, ::rainier_orm::sea_query::Value)>
264 {
265 ::std::vec![ #(#update_vals),* ]
266 }
267 fn pk_value(&self) -> ::rainier_orm::sea_query::Value {
268 ::rainier_orm::ToColumn::to_value(&self.#pk_ident)
269 }
270 }
271 }
272 .into()
273}
274
275fn split_columns(s: &str) -> Vec<String> {
277 s.split(',').map(|c| c.trim().to_string()).filter(|c| !c.is_empty()).collect()
278}
279
280fn parse_references(s: &str) -> (String, String) {
283 let s = s.trim();
284 if let Some(open) = s.find('(') {
285 let table = s[..open].trim().to_string();
286 let col = s[open + 1..].trim_end_matches([')', ' ']).trim().to_string();
287 (table, col)
288 } else if let Some(dot) = s.find('.') {
289 (s[..dot].trim().to_string(), s[dot + 1..].trim().to_string())
290 } else {
291 (s.to_string(), "id".to_string())
292 }
293}
294
295fn ref_action(s: Option<&str>) -> TokenStream2 {
299 let Some(s) = s else {
300 return quote! { ::core::option::Option::None };
301 };
302 let variant = match s.trim().to_ascii_lowercase().replace([' ', '-'], "_").as_str() {
303 "cascade" => quote! { Cascade },
304 "restrict" => quote! { Restrict },
305 "set_null" | "setnull" => quote! { SetNull },
306 "set_default" | "setdefault" => quote! { SetDefault },
307 "no_action" | "noaction" => quote! { NoAction },
308 _ => return quote! { ::core::option::Option::None },
309 };
310 quote! { ::core::option::Option::Some(::rainier_orm::RefAction::#variant) }
311}
312
313fn to_snake_case(s: &str) -> String {
314 let mut out = String::with_capacity(s.len() + 4);
315 for (i, ch) in s.chars().enumerate() {
316 if ch.is_uppercase() {
317 if i != 0 {
318 out.push('_');
319 }
320 out.extend(ch.to_lowercase());
321 } else {
322 out.push(ch);
323 }
324 }
325 out
326}
327
328fn compile_err(msg: &str) -> TokenStream {
329 quote! { compile_error!(#msg); }.into()
330}