rainier_orm_macros/
lib.rs

1//! `#[derive(Entity)]` for rainier_orm.
2//!
3//! From a plain struct it generates the `Entity` impl: table + column
4//! metadata (column SQL type is deferred to the field type's `SqlType`
5//! impl, so it stays open to new types), a row decoder (`from_row`, via
6//! each field type's `FromColumn`), the insert/primary-key value projections
7//! (via `ToColumn`), plus index and foreign-key constraint metadata. No
8//! per-entity SQL is written by hand.
9//!
10//! ```ignore
11//! #[derive(Entity)]
12//! #[orm(table = "posts")]
13//! #[orm(index = "author_id, created_at")]      // composite index
14//! #[orm(unique = "author_id, slug")]           // composite unique
15//! struct Post {
16//!     #[orm(pk, auto_increment)]
17//!     id: u64,
18//!     #[orm(unique)]                            // single-column unique
19//!     slug: String,
20//!     #[orm(index, references = "users(id)", on_delete = "cascade")]
21//!     author_id: u64,                            // FK *id*, not a relationship
22//!     body: Option<String>,                      // nullable inferred from Option
23//!     #[orm(column = "created_at")]
24//!     created: chrono::DateTime<chrono::Utc>,
25//! }
26//! ```
27
28use 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    // --- struct-level: table name + composite index/unique declarations ---
44    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)>, // (foreign_table, foreign_column)
85        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    // --- indexes(): single-column `#[orm(index)]` + struct composites ---
170    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    // --- foreign_keys(): one per `#[orm(references = "...")]` field ---
192    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    // Shard-encoded columns (`#[orm(shard)]`) — routing keys.
211    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    // Inserts skip auto-increment primary keys (let the DB assign them).
224    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    // Updates set every non-primary-key column.
230    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
275/// `"a, b , c"` → `["a", "b", "c"]`.
276fn split_columns(s: &str) -> Vec<String> {
277    s.split(',').map(|c| c.trim().to_string()).filter(|c| !c.is_empty()).collect()
278}
279
280/// Parse `references` targets: `"users(id)"` or `"users.id"` → `("users","id")`.
281/// A bare `"users"` defaults the column to `id`.
282fn 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
295/// Map an `on_delete`/`on_update` string to an `Option<RefAction>` token
296/// stream (`None` when unset; an unknown action becomes `None` and the DB
297/// default applies).
298fn 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}
331
332/// Derive a [test factory] for an entity.
333///
334/// ```ignore
335/// #[derive(Entity, Clone, Default, Factory)]
336/// struct User { id: u64, email: String }
337///
338/// let users = User::factory().count(3).make();
339/// ```
340///
341/// Generates `User::factory()`, building each row from [`Default`] — so the
342/// type needs one, and three rows are three identical rows until a
343/// `.sequence(..)` distinguishes them. Anything with a `UNIQUE` column
344/// therefore needs one; see the trait's own documentation for why that is not
345/// automatic.
346///
347/// The expansion is deliberately tiny. A derive that invented plausible values
348/// per field would have to guess what each column means, and a wrong guess
349/// produces a row that fails to insert for a reason nobody can see from the
350/// test that wrote it.
351///
352/// [test factory]: https://docs.rs/rainier-database/latest/rainier_database/factory/index.html
353#[proc_macro_derive(Factory)]
354pub fn derive_factory(input: TokenStream) -> TokenStream {
355    let input = parse_macro_input!(input as DeriveInput);
356    let ident = input.ident;
357    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
358
359    quote! {
360        impl #impl_generics ::rainier_database::factory::HasFactory for #ident #type_generics
361        #where_clause
362        {
363            fn factory() -> ::rainier_database::factory::Factory<Self> {
364                ::rainier_database::factory::Factory::new(|_index| <Self as Default>::default())
365            }
366        }
367    }
368    .into()
369}
rainier_orm_macros/lib.rs

rainier_orm_macros/
lib.rs
