rdbi_codegen/codegen/
struct_generator.rs

1//! Struct generator - generates Rust structs from table metadata
2
3use crate::error::Result;
4use std::collections::HashSet;
5use std::fs;
6use std::path::Path;
7use tracing::debug;
8
9use crate::config::CodegenConfig;
10use crate::parser::{ColumnMetadata, TableMetadata};
11
12use super::naming::{escape_field_name, to_enum_name, to_enum_variant, to_struct_name};
13use super::type_resolver::TypeResolver;
14
15/// Generate struct files for all tables
16pub fn generate_structs(tables: &[TableMetadata], config: &CodegenConfig) -> Result<()> {
17    let output_dir = &config.output_structs_dir;
18    fs::create_dir_all(output_dir)?;
19
20    // Generate mod.rs with shared pagination types
21    let mut mod_content = String::new();
22    mod_content.push_str("// Generated model structs\n\n");
23
24    for table in tables {
25        let file_name = heck::AsSnakeCase(&table.name).to_string();
26        mod_content.push_str(&format!(
27            "#[allow(dead_code, clippy::all)]\nmod {};\npub use {}::*;\n",
28            file_name, file_name
29        ));
30    }
31
32    mod_content.push('\n');
33    mod_content.push_str(&generate_shared_pagination_types());
34
35    let mod_path = output_dir.join("mod.rs");
36    fs::write(&mod_path, &mod_content)?;
37
38    // Generate each struct file
39    for table in tables {
40        generate_struct_file(table, output_dir)?;
41    }
42
43    Ok(())
44}
45
46/// Generate shared pagination types (SortDirection, PaginatedResult)
47fn generate_shared_pagination_types() -> String {
48    r#"/// Sort direction for pagination
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum SortDirection { Asc, Desc }
51
52impl SortDirection {
53pub fn as_sql(&self) -> &'static str {
54match self { Self::Asc => "ASC", Self::Desc => "DESC" }
55}
56}
57
58/// Paginated result container
59#[derive(Debug, Clone)]
60pub struct PaginatedResult<T> {
61pub items: Vec<T>,
62pub total_count: i64,
63pub current_page: i32,
64pub total_pages: i32,
65pub page_size: i32,
66pub has_next: bool,
67}
68
69impl<T> PaginatedResult<T> {
70pub fn new(items: Vec<T>, total_count: i64, current_page: i32, page_size: i32) -> Self {
71let total_pages = ((total_count as f64) / (page_size as f64)).ceil() as i32;
72let has_next = current_page < total_pages;
73Self { items, total_count, current_page, total_pages, page_size, has_next }
74}
75}
76"#
77    .to_string()
78}
79
80/// Generate a single struct file for a table
81fn generate_struct_file(table: &TableMetadata, output_dir: &Path) -> Result<()> {
82    let struct_name = to_struct_name(&table.name);
83    let file_name = format!("{}.rs", heck::AsSnakeCase(&table.name));
84    debug!("Generating struct {} -> {}", struct_name, file_name);
85
86    let mut code = String::new();
87
88    // Collect enum columns
89    let mut enum_columns: Vec<&ColumnMetadata> = Vec::new();
90    for col in &table.columns {
91        if col.is_enum() {
92            enum_columns.push(col);
93        }
94    }
95
96    // Generate imports (serde and rdbi)
97    code.push_str("use serde::{Deserialize, Serialize};\n");
98    code.push('\n');
99
100    // Generate enum types first
101    for col in &enum_columns {
102        if let Some(values) = &col.enum_values {
103            code.push_str(&generate_enum(&table.name, col, values));
104            code.push('\n');
105        }
106    }
107
108    // Generate struct documentation
109    code.push_str(&format!("/// Database table: `{}`\n", table.name));
110    if let Some(comment) = &table.comment {
111        if !comment.is_empty() {
112            code.push_str(&format!("///\n/// {}\n", comment));
113        }
114    }
115
116    // Generate struct with derives
117    code.push_str("#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, rdbi::FromRow, rdbi::ToParams)]\n");
118    code.push_str(&format!("pub struct {} {{\n", struct_name));
119
120    // Generate fields
121    for col in &table.columns {
122        let field_name = escape_field_name(&col.name);
123        let rust_type = TypeResolver::resolve(col, &table.name);
124
125        // Add field documentation
126        code.push_str(&format!("/// Column: `{}`", col.name));
127        let index_info = get_index_info(table, &col.name);
128        if !index_info.is_empty() {
129            code.push_str(&format!(" ({})", index_info.join(", ")));
130        }
131        if let Some(comment) = &col.comment {
132            if !comment.is_empty() {
133                code.push_str(&format!(" - {}", comment));
134            }
135        }
136        code.push('\n');
137
138        // Add rdbi attributes
139        let mut attrs = Vec::new();
140        if field_name != col.name {
141            attrs.push(format!("rename = \"{}\"", col.name));
142        }
143        if col.is_auto_increment {
144            attrs.push("skip_insert".to_string());
145        }
146        if !attrs.is_empty() {
147            code.push_str(&format!("#[rdbi({})]\n", attrs.join(", ")));
148        }
149
150        // Add serde rename for raw identifiers (r#type -> "type")
151        if field_name != col.name {
152            code.push_str(&format!("#[serde(rename = \"{}\")]\n", col.name));
153        }
154
155        code.push_str(&format!(
156            "pub {}: {},\n",
157            field_name,
158            rust_type.to_type_string()
159        ));
160    }
161
162    code.push_str("}\n");
163    code.push('\n');
164    code.push_str(&generate_sort_by_enum(table));
165
166    let file_path = output_dir.join(&file_name);
167    fs::write(&file_path, &code)?;
168    Ok(())
169}
170
171/// Generate SortBy enum for a table (used in pagination)
172fn generate_sort_by_enum(table: &TableMetadata) -> String {
173    let struct_name = to_struct_name(&table.name);
174    let enum_name = format!("{}SortBy", struct_name);
175
176    let mut code = String::new();
177
178    code.push_str(&format!("/// Sort columns for `{}`\n", table.name));
179    code.push_str("#[derive(Debug, Clone, Copy, PartialEq, Eq)]\n");
180    code.push_str(&format!("pub enum {} {{\n", enum_name));
181    for col in &table.columns {
182        let variant = heck::AsPascalCase(&col.name).to_string();
183        code.push_str(&format!("{},\n", variant));
184    }
185    code.push_str("}\n\n");
186
187    code.push_str(&format!("impl {} {{\n", enum_name));
188    code.push_str("pub fn as_sql(&self) -> &'static str {\n");
189    code.push_str("match self {\n");
190    for col in &table.columns {
191        let variant = heck::AsPascalCase(&col.name).to_string();
192        code.push_str(&format!("Self::{} => \"`{}`\",\n", variant, col.name));
193    }
194    code.push_str("}\n}\n}\n");
195
196    code
197}
198
199/// Generate an enum type for a column
200fn generate_enum(table_name: &str, column: &ColumnMetadata, values: &[String]) -> String {
201    let enum_name = to_enum_name(table_name, &column.name);
202    let mut code = String::new();
203
204    code.push_str(&format!("/// Enum for `{}.{}`\n", table_name, column.name));
205    code.push_str("#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]\n");
206    code.push_str(&format!("pub enum {} {{\n", enum_name));
207
208    let mut used_variants: HashSet<String> = HashSet::new();
209    let mut variant_mappings: Vec<(String, String)> = Vec::new();
210
211    for value in values {
212        let variant = to_enum_variant(value);
213        let final_variant = if used_variants.contains(&variant) {
214            let mut counter = 2;
215            loop {
216                let new_variant = format!("{}{}", variant, counter);
217                if !used_variants.contains(&new_variant) {
218                    break new_variant;
219                }
220                counter += 1;
221            }
222        } else {
223            variant
224        };
225        used_variants.insert(final_variant.clone());
226
227        let clean_value = value.trim_matches('\'').trim_matches('"');
228        if final_variant != clean_value {
229            code.push_str(&format!("#[serde(rename = \"{}\")]\n", clean_value));
230        }
231        code.push_str(&format!("{},\n", final_variant));
232        variant_mappings.push((final_variant, clean_value.to_string()));
233    }
234
235    code.push_str("}\n\n");
236
237    // FromValue implementation
238    code.push_str(&format!("impl rdbi::FromValue for {} {{\n", enum_name));
239    code.push_str("fn from_value(value: rdbi::Value) -> rdbi::Result<Self> {\n");
240    code.push_str("match value {\n");
241    code.push_str("rdbi::Value::String(s) => match s.as_str() {\n");
242    for (variant, db_value) in &variant_mappings {
243        code.push_str(&format!("\"{}\" => Ok(Self::{}),\n", db_value, variant));
244    }
245    code.push_str(&format!(
246        "_ => Err(rdbi::Error::TypeConversion {{ expected: \"{}\", actual: s }}),\n",
247        enum_name
248    ));
249    code.push_str("},\n");
250    code.push_str(&format!(
251        "_ => Err(rdbi::Error::TypeConversion {{ expected: \"{}\", actual: value.type_name().to_string() }}),\n",
252        enum_name
253    ));
254    code.push_str("}\n}\n}\n\n");
255
256    // ToValue implementation
257    code.push_str(&format!("impl rdbi::ToValue for {} {{\n", enum_name));
258    code.push_str("fn to_value(&self) -> rdbi::Value {\n");
259    code.push_str("rdbi::Value::String(match self {\n");
260    for (variant, db_value) in &variant_mappings {
261        code.push_str(&format!(
262            "Self::{} => \"{}\".to_string(),\n",
263            variant, db_value
264        ));
265    }
266    code.push_str("})\n}\n}\n");
267
268    code
269}
270
271/// Get index information for a column
272fn get_index_info(table: &TableMetadata, column_name: &str) -> Vec<String> {
273    let mut info = Vec::new();
274
275    // Check primary key
276    if let Some(pk) = &table.primary_key {
277        if pk.columns.contains(&column_name.to_string()) {
278            info.push("PRIMARY KEY".to_string());
279        }
280    }
281
282    // Check indexes
283    for index in &table.indexes {
284        if index.columns.contains(&column_name.to_string()) {
285            let label = if index.unique {
286                format!("UNIQUE: {}", index.name)
287            } else {
288                format!("INDEX: {}", index.name)
289            };
290            info.push(label);
291        }
292    }
293
294    info
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use crate::parser::{IndexMetadata, PrimaryKey};
301
302    fn make_table() -> TableMetadata {
303        TableMetadata {
304            name: "users".to_string(),
305            comment: None,
306            columns: vec![
307                ColumnMetadata {
308                    name: "id".to_string(),
309                    data_type: "BIGINT".to_string(),
310                    nullable: false,
311                    default_value: None,
312                    is_auto_increment: true,
313                    is_unsigned: false,
314                    enum_values: None,
315                    comment: None,
316                },
317                ColumnMetadata {
318                    name: "username".to_string(),
319                    data_type: "VARCHAR(255)".to_string(),
320                    nullable: false,
321                    default_value: None,
322                    is_auto_increment: false,
323                    is_unsigned: false,
324                    enum_values: None,
325                    comment: None,
326                },
327                ColumnMetadata {
328                    name: "status".to_string(),
329                    data_type: "ENUM".to_string(),
330                    nullable: false,
331                    default_value: None,
332                    is_auto_increment: false,
333                    is_unsigned: false,
334                    enum_values: Some(vec![
335                        "ACTIVE".to_string(),
336                        "INACTIVE".to_string(),
337                        "PENDING".to_string(),
338                    ]),
339                    comment: None,
340                },
341            ],
342            indexes: vec![IndexMetadata {
343                name: "idx_username".to_string(),
344                columns: vec!["username".to_string()],
345                unique: true,
346            }],
347            foreign_keys: vec![],
348            primary_key: Some(PrimaryKey {
349                columns: vec!["id".to_string()],
350            }),
351        }
352    }
353
354    #[test]
355    fn test_get_index_info() {
356        let table = make_table();
357        let info = get_index_info(&table, "id");
358        assert!(info.contains(&"PRIMARY KEY".to_string()));
359
360        let info = get_index_info(&table, "username");
361        assert!(info.iter().any(|i| i.contains("UNIQUE")));
362    }
363
364    #[test]
365    fn test_generate_enum() {
366        let col = ColumnMetadata {
367            name: "status".to_string(),
368            data_type: "ENUM".to_string(),
369            nullable: false,
370            default_value: None,
371            is_auto_increment: false,
372            is_unsigned: false,
373            enum_values: Some(vec!["ACTIVE".to_string(), "INACTIVE".to_string()]),
374            comment: None,
375        };
376
377        let code = generate_enum("users", &col, col.enum_values.as_ref().unwrap());
378        assert!(code.contains("pub enum UsersStatus"));
379        assert!(code.contains("Active"));
380        assert!(code.contains("Inactive"));
381        // Check for rdbi trait implementations
382        assert!(code.contains("impl rdbi::FromValue for UsersStatus"));
383        assert!(code.contains("impl rdbi::ToValue for UsersStatus"));
384    }
385}
rdbi_codegen/codegen/struct_generator.rs

rdbi_codegen/codegen/
struct_generator.rs
