migrations_macro_rs/
lib.rs

1use quote::quote;
2use syn::__private::TokenStream;
3use syn::{Data, DeriveInput, Fields, Meta, parse_macro_input};
4
5#[proc_macro_derive(Csv, attributes(migrations))]
6pub fn csv_macro(input: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(input as DeriveInput);
8    let struct_name = input.ident;
9    let expanded = quote! {
10        impl #struct_name {
11            /// Save the data to a CSV file.
12            pub async fn save_csv(
13                records: Vec<Self>,
14                file_path: &str,
15            ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
16            where
17                Self: serde::Serialize,
18            {
19                let mut writer = csv::Writer::from_path(file_path)?;
20                for record in records {
21                    writer.serialize(record)?;
22                }
23                writer.flush()?;
24                Ok(())
25            }
26            /// Load data from a CSV file.
27            pub async fn load_csv(
28                file_path: &str,
29            ) -> Result<Vec<Self>, Box<dyn std::error::Error + Send + Sync>>
30            where
31                Self: for<'de> serde::Deserialize<'de>,
32            {
33                let mut reader = csv::Reader::from_path(file_path)?;
34                let mut records = Vec::new();
35                for result in reader.deserialize() {
36                    let record: Self = result?;
37                    records.push(record);
38                }
39                Ok(records)
40            }
41
42            /// Save field values of vec self by field_extractor into csv
43            pub async fn save_field<F, T>(
44                records: Vec<Self>,
45                file_path: &str,
46                field_extractor: F,
47            ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
48            where
49                F: Fn(&Self) -> T,
50                T: serde::Serialize,
51            {
52                let mut writer = csv::Writer::from_path(file_path)?;
53                for record in records {
54                    let field_value = field_extractor(&record);
55                    writer.serialize(field_value)?;
56                }
57                writer.flush()?;
58                Ok(())
59            }
60        }
61    };
62    TokenStream::from(expanded)
63}
64
65#[proc_macro_derive(Persistency, attributes(migrations))]
66pub fn persistency_macro(input: TokenStream) -> TokenStream {
67    let input = parse_macro_input!(input as DeriveInput);
68    let struct_name = input.ident;
69
70    let mut schema = String::new();
71    let mut table = String::new();
72
73    for attr in input.attrs {
74        if attr.path().is_ident("migrations") {
75            match attr.meta {
76                Meta::List(list) => {
77                    let nested = list
78                        .parse_args_with(
79                            syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated,
80                        )
81                        .unwrap_or_else(|_| syn::punctuated::Punctuated::new());
82
83                    for meta in nested {
84                        match meta {
85                            Meta::NameValue(nv) if nv.path.is_ident("schema") => {
86                                if let syn::Expr::Lit(syn::ExprLit {
87                                    lit: syn::Lit::Str(lit_str),
88                                    ..
89                                }) = nv.value
90                                {
91                                    schema = lit_str.value();
92                                }
93                            }
94                            Meta::NameValue(nv) if nv.path.is_ident("table") => {
95                                if let syn::Expr::Lit(syn::ExprLit {
96                                    lit: syn::Lit::Str(lit_str),
97                                    ..
98                                }) = nv.value
99                                {
100                                    table = lit_str.value();
101                                }
102                            }
103                            _ => {}
104                        }
105                    }
106                }
107                _ => {}
108            }
109        }
110    }
111
112    if schema.is_empty() || table.is_empty() {
113        return TokenStream::from(quote! {
114            compile_error!("Both `schema` and `table` must be provided as attributes for #[migrations(schema = \"...\", table = \"...\")].");
115        });
116    }
117
118    let expanded = quote! {
119        impl #struct_name {
120            /// Sets the table to LOGGED (persistent).
121            pub async fn set_persistent(
122                pool: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
123            ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
124                let client = pool.get().await?;
125                let query = format!(
126                    r#"ALTER TABLE {}."{}" SET LOGGED;"#, #schema, #table
127                );
128                client.execute(&query, &[]).await?;
129                Ok(())
130            }
131
132            /// Sets the table to UNLOGGED (not persistent).
133            pub async fn set_unlogged(
134                pool: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
135            ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
136                let client = pool.get().await?;
137                let query = format!(
138                    r#"ALTER TABLE {}."{}" SET UNLOGGED;"#, #schema, #table
139                );
140                client.execute(&query, &[]).await?;
141                Ok(())
142            }
143        }
144    };
145
146    TokenStream::from(expanded)
147}
148
149#[proc_macro_derive(Vacuum, attributes(migrations))]
150pub fn vacuum_macro(input: TokenStream) -> TokenStream {
151    let input = parse_macro_input!(input as DeriveInput);
152    let struct_name = input.ident;
153
154    let mut schema = String::new();
155    let mut table = String::new();
156
157    for attr in input.attrs {
158        if attr.path().is_ident("migrations") {
159            match attr.meta {
160                Meta::List(list) => {
161                    let nested = list
162                        .parse_args_with(
163                            syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated,
164                        )
165                        .unwrap_or_else(|_| syn::punctuated::Punctuated::new());
166
167                    for meta in nested {
168                        match meta {
169                            Meta::NameValue(nv) if nv.path.is_ident("schema") => {
170                                if let syn::Expr::Lit(syn::ExprLit {
171                                    lit: syn::Lit::Str(lit_str),
172                                    ..
173                                }) = nv.value
174                                {
175                                    schema = lit_str.value();
176                                }
177                            }
178                            Meta::NameValue(nv) if nv.path.is_ident("table") => {
179                                if let syn::Expr::Lit(syn::ExprLit {
180                                    lit: syn::Lit::Str(lit_str),
181                                    ..
182                                }) = nv.value
183                                {
184                                    table = lit_str.value();
185                                }
186                            }
187                            _ => {}
188                        }
189                    }
190                }
191                _ => {}
192            }
193        }
194    }
195
196    if schema.is_empty() || table.is_empty() {
197        return TokenStream::from(quote! {
198            compile_error!("Both `schema` and `table` must be provided as attributes for #[migrations(schema = \"...\", table = \"...\")].");
199        });
200    }
201
202    let expanded = quote! {
203        impl #struct_name {
204            /// Executes a VACUUM FULL on the table.
205            pub async fn vacuum(
206                pool: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
207            ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
208                let client = pool.get().await?;
209                let query = format!(r#"VACUUM FULL {}."{}";"#, #schema, #table);
210                client.execute(&query, &[]).await?;
211                Ok(())
212            }
213        }
214    };
215
216    TokenStream::from(expanded)
217}
218
219#[proc_macro_derive(Reindex, attributes(migrations))]
220pub fn reindex_macro(input: TokenStream) -> TokenStream {
221    let input = parse_macro_input!(input as DeriveInput);
222    let struct_name = input.ident;
223
224    let mut schema = String::new();
225    let mut table = String::new();
226
227    for attr in input.attrs {
228        if attr.path().is_ident("migrations") {
229            match attr.meta {
230                Meta::List(list) => {
231                    let nested = list
232                        .parse_args_with(
233                            syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated,
234                        )
235                        .unwrap_or_else(|_| syn::punctuated::Punctuated::new());
236
237                    for meta in nested {
238                        match meta {
239                            Meta::NameValue(nv) if nv.path.is_ident("schema") => {
240                                if let syn::Expr::Lit(syn::ExprLit {
241                                    lit: syn::Lit::Str(lit_str),
242                                    ..
243                                }) = nv.value
244                                {
245                                    schema = lit_str.value();
246                                }
247                            }
248                            Meta::NameValue(nv) if nv.path.is_ident("table") => {
249                                if let syn::Expr::Lit(syn::ExprLit {
250                                    lit: syn::Lit::Str(lit_str),
251                                    ..
252                                }) = nv.value
253                                {
254                                    table = lit_str.value();
255                                }
256                            }
257                            _ => {}
258                        }
259                    }
260                }
261                _ => {}
262            }
263        }
264    }
265
266    if schema.is_empty() || table.is_empty() {
267        return TokenStream::from(quote! {
268            compile_error!("Both `schema` and `table` must be provided as attributes for #[migrations(schema = \"...\", table = \"...\")].");
269        });
270    }
271
272    let expanded = quote! {
273        impl #struct_name {
274            /// Reindexes the table.
275            pub async fn reindex(
276                pool: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
277            ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
278                let client = pool.get().await?;
279                let query = format!(r#"REINDEX TABLE {}."{}";"#, #schema, #table);
280                client.execute(&query, &[]).await?;
281                Ok(())
282            }
283        }
284    };
285
286    TokenStream::from(expanded)
287}
288
289#[proc_macro_derive(Database, attributes(migrations))]
290pub fn database_macro(input: TokenStream) -> TokenStream {
291    let input = parse_macro_input!(input as DeriveInput);
292    let struct_name = input.ident;
293
294    // Extract schema and table from attributes
295    let mut schema = String::new();
296    let mut table = String::new();
297
298    for attr in input.attrs.iter() {
299        if attr.path().is_ident("migrations") {
300            match attr.meta {
301                Meta::List(ref list) => {
302                    let nested = list
303                        .parse_args_with(
304                            syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated,
305                        )
306                        .unwrap_or_else(|_| syn::punctuated::Punctuated::new());
307
308                    for meta in nested {
309                        match meta {
310                            Meta::NameValue(ref nv) if nv.path.is_ident("schema") => {
311                                if let syn::Expr::Lit(syn::ExprLit {
312                                    lit: syn::Lit::Str(ref lit_str),
313                                    ..
314                                }) = nv.value
315                                {
316                                    schema = lit_str.value();
317                                }
318                            }
319                            Meta::NameValue(ref nv) if nv.path.is_ident("table") => {
320                                if let syn::Expr::Lit(syn::ExprLit {
321                                    lit: syn::Lit::Str(ref lit_str),
322                                    ..
323                                }) = nv.value
324                                {
325                                    table = lit_str.value();
326                                }
327                            }
328                            _ => {}
329                        }
330                    }
331                }
332                _ => {}
333            }
334        }
335    }
336
337    if schema.is_empty() || table.is_empty() {
338        return TokenStream::from(quote! {
339            compile_error!("Both `schema` and `table` must be provided as attributes for #[migrations(schema = \"...\", table = \"...\")].");
340        });
341    }
342
343    // Extract field names and their serde renames
344    let mut field_conversions = Vec::new();
345    let mut field_names = Vec::new();
346
347    if let Data::Struct(data_struct) = &input.data {
348        if let Fields::Named(fields_named) = &data_struct.fields {
349            for field in fields_named.named.iter() {
350                if let Some(field_name) = &field.ident {
351                    let mut db_field_name = field_name.to_string();
352
353                    // Look for serde rename attribute
354                    for attr in &field.attrs {
355                        if attr.path().is_ident("serde") {
356                            if let Meta::List(list) = &attr.meta {
357                                let nested = list.parse_args_with(syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated)
358                                    .unwrap_or_else(|_| syn::punctuated::Punctuated::new());
359
360                                for meta in nested {
361                                    if let Meta::NameValue(nv) = meta {
362                                        if nv.path.is_ident("rename") {
363                                            if let syn::Expr::Lit(syn::ExprLit {
364                                                lit: syn::Lit::Str(lit_str),
365                                                ..
366                                            }) = nv.value
367                                            {
368                                                db_field_name = lit_str.value();
369                                            }
370                                        }
371                                    }
372                                }
373                            }
374                        }
375                    }
376
377                    // Add to field conversions and names
378                    let field_name_ident = field_name;
379                    let db_field_name_literal = db_field_name.clone();
380
381                    field_conversions.push(quote! {
382                        #field_name_ident: row.get(#db_field_name_literal)
383                    });
384
385                    field_names.push(db_field_name);
386                }
387            }
388        }
389    }
390
391    // Prepare comma-separated list of field names for SQL query
392    let field_names_str = field_names
393        .iter()
394        .map(|name| format!("\"{}\"", name))
395        .collect::<Vec<_>>()
396        .join(", ");
397
398    let expanded = quote! {
399        impl #struct_name {
400            /// Converts a PostgreSQL Row to a struct
401            pub fn convert_from_row(row: bb8_postgres::tokio_postgres::Row) -> Self {
402                Self {
403                    #(#field_conversions),*
404                }
405            }
406
407            /// Fetches all records from the table
408            pub async fn get(
409                connection: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
410            ) -> Result<Vec<Self>, Box<dyn std::error::Error + Send + Sync>> {
411                let client = connection.get().await?;
412                let query = format!(
413                    r#"SELECT {} FROM {}."{}"#,
414                    #field_names_str, #schema, #table
415                );
416                let rows = client.query(&query, &[]).await?;
417                let results: Vec<Self> = rows.into_iter().map(Self::convert_from_row).collect();
418                Ok(results)
419            }
420
421            /// Fetches records and filters them by a specified i64 field
422            pub async fn get_filter_i64<F>(
423                connection: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
424                field_extractor: F,
425                filter_values: std::collections::HashSet<i64>,
426            ) -> Result<Vec<Self>, Box<dyn std::error::Error + Send + Sync>>
427            where
428                F: Fn(&Self) -> i64,
429            {
430                let client = connection.get().await?;
431                let query = format!(
432                    r#"SELECT {} FROM {}."{}"#,
433                    #field_names_str, #schema, #table
434                );
435                let rows = client.query(&query, &[]).await?;
436                let results: Vec<Self> = rows
437                    .into_iter()
438                    .map(Self::convert_from_row)
439                    .filter(|record| filter_values.contains(&field_extractor(record)))
440                    .collect();
441
442                Ok(results)
443            }
444
445            /// Fetches records and filters them by two specified i64 fields
446            pub async fn get_filter_i64_i64<F1, F2>(
447                connection: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
448                field_extractor1: F1,
449                filter_values1: std::collections::HashSet<i64>,
450                field_extractor2: F2,
451                filter_values2: std::collections::HashSet<i64>,
452            ) -> Result<Vec<Self>, Box<dyn std::error::Error + Send + Sync>>
453            where
454                F1: Fn(&Self) -> Option<i64>,
455                F2: Fn(&Self) -> Option<i64>,
456            {
457                let client = connection.get().await?;
458                let query = format!(
459                    r#"SELECT {} FROM {}."{}"#,
460                    #field_names_str, #schema, #table
461                );
462                let rows = client.query(&query, &[]).await?;
463                let results: Vec<Self> = rows
464                    .into_iter()
465                    .map(Self::convert_from_row)
466                    .filter(|record| {
467                        let value1 = field_extractor1(record).unwrap_or(-1);
468                        let value2 = field_extractor2(record).unwrap_or(-1);
469                        filter_values1.contains(&value1) && filter_values2.contains(&value2)
470                    })
471                    .collect();
472
473                Ok(results)
474            }
475
476            /// Fetches records and filters them by a specified uuid field
477            pub async fn get_filter_uuid<F>(
478                connection: &bb8_postgres::bb8::Pool<bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>>,
479                field_extractor: F,
480                filter_values: std::collections::HashSet<uuid>,
481            ) -> Result<Vec<Self>, Box<dyn std::error::Error + Send + Sync>>
482            where
483                F: Fn(&Self) -> uuid,
484            {
485                let client = connection.get().await?;
486                let query = format!(
487                    r#"SELECT {} FROM {}."{}"#,
488                    #field_names_str, #schema, #table
489                );
490                let rows = client.query(&query, &[]).await?;
491                let results: Vec<Self> = rows
492                    .into_iter()
493                    .map(Self::convert_from_row)
494                    .filter(|record| filter_values.contains(&field_extractor(record)))
495                    .collect();
496
497                Ok(results)
498            }
499        }
500    };
501
502    TokenStream::from(expanded)
503}
migrations_macro_rs/lib.rs

migrations_macro_rs/
lib.rs
