sz_orm_core/
schema_gen.rs

1//! Diesel 风格 schema.rs 自动生成
2//!
3//! 从数据库表元数据生成 `typed_query!` 表声明，配合宏做编译期列名校验。
4//!
5//! # 设计
6//!
7//! Diesel 的 `diesel print-schema` 命令从数据库反向生成 `schema.rs` 文件，
8//! 包含 `table!` 宏声明，让 SQL 列名错误在编译期被捕获。
9//!
10//! 本模块提供等价功能，生成 SZ-ORM 的 `typed_query!` 声明：
11//!
12//! ```ignore
13//! // 生成的 schema.rs 内容
14//! use sz_orm_core::typed_query;
15//!
16//! typed_query! {
17//!     table users {
18//!         id: i64,
19//!         name: String,
20//!         email: String,
21//!         created_at: String,
22//!     }
23//! }
24//!
25//! typed_query! {
26//!     table orders {
27//!         order_id: i64,
28//!         user_id: i64,
29//!         total: f64,
30//!     }
31//! }
32//! ```
33
34use std::fmt::Write as _;
35
36/// 单列的元数据（足够生成 typed_query! 声明）
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct ColumnSchema {
39    /// 列名
40    pub name: String,
41    /// Rust 类型名（如 "i64"、"String"、"f64"、"`Option<String>`"）
42    pub rust_type: String,
43}
44
45/// 单表的元数据
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct TableSchema {
48    /// 表名
49    pub name: String,
50    /// 所有列
51    pub columns: Vec<ColumnSchema>,
52}
53
54/// 生成器：把表元数据列表转换成 schema.rs 文件内容
55pub struct SchemaGenerator {
56    /// 文件头注释（默认包含 SZ-ORM 自动生成提示）
57    header: String,
58    /// 是否生成 `use` 语句
59    emit_use: bool,
60}
61
62impl Default for SchemaGenerator {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68impl SchemaGenerator {
69    /// 创建新的生成器
70    pub fn new() -> Self {
71        Self {
72            header: format!(
73                "// Auto-generated by sz-orm-cli generate schema at {}\n\
74                 // DO NOT EDIT MANUALLY — re-run the command to refresh.\n\
75                 //\n\
76                 // This file contains typed_query! table declarations\n\
77                 // enabling compile-time column-name verification.\n",
78                chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC")
79            ),
80            emit_use: true,
81        }
82    }
83
84    /// 自定义文件头注释
85    pub fn with_header(mut self, header: impl Into<String>) -> Self {
86        self.header = header.into();
87        self
88    }
89
90    /// 是否生成 `use sz_orm_core::typed_query;` 语句
91    pub fn emit_use(mut self, emit: bool) -> Self {
92        self.emit_use = emit;
93        self
94    }
95
96    /// 生成完整的 schema.rs 文件内容
97    pub fn generate(&self, tables: &[TableSchema]) -> String {
98        let mut out = String::new();
99
100        // 文件头
101        writeln!(out, "{}", self.header).unwrap();
102        writeln!(out).unwrap();
103
104        // use 语句
105        if self.emit_use {
106            writeln!(out, "use sz_orm_core::typed_query;").unwrap();
107            writeln!(out).unwrap();
108        }
109
110        // 每张表生成一个 typed_query! 声明
111        for (idx, table) in tables.iter().enumerate() {
112            if idx > 0 {
113                writeln!(out).unwrap();
114            }
115            write!(out, "{}", self.render_table(table)).unwrap();
116        }
117
118        out
119    }
120
121    /// 渲染单张表的 typed_query! 声明
122    fn render_table(&self, table: &TableSchema) -> String {
123        let mut out = String::new();
124        writeln!(out, "typed_query! {{").unwrap();
125        writeln!(out, "    table {} {{", table.name).unwrap();
126        for col in &table.columns {
127            writeln!(out, "        {}: {},", col.name, col.rust_type).unwrap();
128        }
129        writeln!(out, "    }}").unwrap();
130        writeln!(out, "}}").unwrap();
131        out
132    }
133}
134
135/// 把 SQL 类型字符串映射到 Rust 类型字符串
136///
137/// 用于 `generate schema` 命令从 DB 元数据生成 typed_query! 声明
138pub fn sql_type_to_rust(sql_type: &str, nullable: bool) -> String {
139    let base = match sql_type.to_lowercase() {
140        // 整数（按长度优先匹配，避免 "int" 误匹配 "bigint"）
141        s if s.contains("tinyint") => "i8",
142        s if s.contains("smallint") || s.contains("int2") => "i16",
143        s if s.contains("bigint") || s.contains("int8") => "i64",
144        s if s.contains("int") || s.contains("serial") => "i32",
145        // 浮点（更具体的先匹配：float8/float4 在 float 之前）
146        s if s.contains("float8") || s.contains("double") => "f64",
147        s if s.contains("float4") || s.contains("real") => "f32",
148        s if s.contains("float") => "f32",
149        s if s.contains("decimal") || s.contains("numeric") => "f64",
150        // 布尔
151        s if s.contains("bool") => "bool",
152        // 字节
153        s if s.contains("blob") || s.contains("bytea") || s.contains("binary") => "Vec<u8>",
154        // 时间
155        s if s.contains("date") && !s.contains("datetime") && !s.contains("timestamp") => "String",
156        s if s.contains("datetime") || s.contains("timestamp") => "String",
157        s if s.contains("time") => "String",
158        // JSON
159        s if s.contains("json") => "String",
160        // UUID
161        s if s.contains("uuid") => "String",
162        // 字符串
163        s if s.contains("char") || s.contains("text") || s.contains("varchar") => "String",
164        // 默认
165        _ => "String",
166    };
167
168    if nullable {
169        format!("Option<{}>", base)
170    } else {
171        base.to_string()
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn test_sql_type_to_rust_int() {
181        assert_eq!(sql_type_to_rust("INT", false), "i32");
182        assert_eq!(sql_type_to_rust("BIGINT", false), "i64");
183        assert_eq!(sql_type_to_rust("SMALLINT", false), "i16");
184        assert_eq!(sql_type_to_rust("TINYINT", false), "i8");
185    }
186
187    #[test]
188    fn test_sql_type_to_rust_float() {
189        assert_eq!(sql_type_to_rust("FLOAT", false), "f32");
190        assert_eq!(sql_type_to_rust("DOUBLE", false), "f64");
191        assert_eq!(sql_type_to_rust("DECIMAL(10,2)", false), "f64");
192    }
193
194    #[test]
195    fn test_sql_type_to_rust_bool() {
196        assert_eq!(sql_type_to_rust("BOOLEAN", false), "bool");
197        assert_eq!(sql_type_to_rust("TINYINT(1)", false), "i8");
198    }
199
200    #[test]
201    fn test_sql_type_to_rust_string() {
202        assert_eq!(sql_type_to_rust("VARCHAR(255)", false), "String");
203        assert_eq!(sql_type_to_rust("TEXT", false), "String");
204        assert_eq!(sql_type_to_rust("CHAR(36)", false), "String");
205    }
206
207    #[test]
208    fn test_sql_type_to_rust_binary() {
209        assert_eq!(sql_type_to_rust("BLOB", false), "Vec<u8>");
210        assert_eq!(sql_type_to_rust("BYTEA", false), "Vec<u8>");
211    }
212
213    #[test]
214    fn test_sql_type_to_rust_nullable() {
215        assert_eq!(sql_type_to_rust("INT", true), "Option<i32>");
216        assert_eq!(sql_type_to_rust("VARCHAR(255)", true), "Option<String>");
217    }
218
219    #[test]
220    fn test_sql_type_to_rust_pg_types() {
221        assert_eq!(sql_type_to_rust("int8", false), "i64");
222        assert_eq!(sql_type_to_rust("int2", false), "i16");
223        assert_eq!(sql_type_to_rust("float8", false), "f64");
224        assert_eq!(sql_type_to_rust("float4", false), "f32");
225        assert_eq!(sql_type_to_rust("bytea", false), "Vec<u8>");
226    }
227
228    #[test]
229    fn test_sql_type_to_rust_json_uuid() {
230        assert_eq!(sql_type_to_rust("JSON", false), "String");
231        assert_eq!(sql_type_to_rust("JSONB", false), "String");
232        assert_eq!(sql_type_to_rust("UUID", false), "String");
233    }
234
235    #[test]
236    fn test_schema_generator_single_table() {
237        let gen = SchemaGenerator::new().emit_use(false);
238        let tables = vec![TableSchema {
239            name: "users".to_string(),
240            columns: vec![
241                ColumnSchema {
242                    name: "id".to_string(),
243                    rust_type: "i64".to_string(),
244                },
245                ColumnSchema {
246                    name: "name".to_string(),
247                    rust_type: "String".to_string(),
248                },
249            ],
250        }];
251        let output = gen.generate(&tables);
252
253        assert!(output.contains("table users {"));
254        assert!(output.contains("id: i64,"));
255        assert!(output.contains("name: String,"));
256        assert!(output.contains("typed_query! {"));
257    }
258
259    #[test]
260    fn test_schema_generator_multiple_tables() {
261        let gen = SchemaGenerator::new().emit_use(false);
262        let tables = vec![
263            TableSchema {
264                name: "users".to_string(),
265                columns: vec![ColumnSchema {
266                    name: "id".to_string(),
267                    rust_type: "i64".to_string(),
268                }],
269            },
270            TableSchema {
271                name: "orders".to_string(),
272                columns: vec![ColumnSchema {
273                    name: "order_id".to_string(),
274                    rust_type: "i64".to_string(),
275                }],
276            },
277        ];
278        let output = gen.generate(&tables);
279
280        assert!(output.contains("table users {"));
281        assert!(output.contains("table orders {"));
282        // 两个表声明之间应有空行
283        let users_end = output.find("}").unwrap();
284        let orders_start = output.find("table orders").unwrap();
285        let between = &output[users_end..orders_start];
286        assert!(between.contains("\n\n"));
287    }
288
289    #[test]
290    fn test_schema_generator_with_use_statement() {
291        let gen = SchemaGenerator::new().emit_use(true);
292        let tables = vec![TableSchema {
293            name: "t".to_string(),
294            columns: vec![],
295        }];
296        let output = gen.generate(&tables);
297
298        assert!(output.contains("use sz_orm_core::typed_query;"));
299    }
300
301    #[test]
302    fn test_schema_generator_header() {
303        let gen = SchemaGenerator::new();
304        let output = gen.generate(&[]);
305        assert!(output.contains("Auto-generated"));
306        assert!(output.contains("DO NOT EDIT MANUALLY"));
307    }
308
309    #[test]
310    fn test_schema_generator_custom_header() {
311        let gen = SchemaGenerator::new().with_header("// Custom header\n");
312        let output = gen.generate(&[]);
313        assert!(output.starts_with("// Custom header"));
314        assert!(!output.contains("Auto-generated"));
315    }
316
317    #[test]
318    fn test_schema_generator_empty_tables() {
319        let gen = SchemaGenerator::new().emit_use(false);
320        let output = gen.generate(&[]);
321        // 空表列表也应生成（仅有 header，不应有 typed_query! 块）
322        assert!(output.contains("Auto-generated"));
323        // 应该没有 typed_query! 块（header 里只是描述性文字，不算块）
324        // 通过检查 "typed_query! {" 来判断是否有实际声明
325        assert!(!output.contains("typed_query! {"));
326    }
327
328    #[test]
329    fn test_schema_generator_option_type() {
330        let gen = SchemaGenerator::new().emit_use(false);
331        let tables = vec![TableSchema {
332            name: "products".to_string(),
333            columns: vec![ColumnSchema {
334                name: "price".to_string(),
335                rust_type: "Option<f64>".to_string(),
336            }],
337        }];
338        let output = gen.generate(&tables);
339
340        assert!(output.contains("price: Option<f64>,"));
341    }
342
343    #[test]
344    fn test_schema_generator_compound_type() {
345        let gen = SchemaGenerator::new().emit_use(false);
346        let tables = vec![TableSchema {
347            name: "files".to_string(),
348            columns: vec![ColumnSchema {
349                name: "content".to_string(),
350                rust_type: "Vec<u8>".to_string(),
351            }],
352        }];
353        let output = gen.generate(&tables);
354
355        assert!(output.contains("content: Vec<u8>,"));
356    }
357
358    #[test]
359    fn test_generated_code_is_valid_syntax() {
360        // 验证生成的代码包含正确的 typed_query! 调用语法
361        let gen = SchemaGenerator::new().emit_use(false);
362        let tables = vec![TableSchema {
363            name: "typed_validate_test".to_string(),
364            columns: vec![
365                ColumnSchema {
366                    name: "id".to_string(),
367                    rust_type: "i64".to_string(),
368                },
369                ColumnSchema {
370                    name: "name".to_string(),
371                    rust_type: "String".to_string(),
372                },
373            ],
374        }];
375        let output = gen.generate(&tables);
376
377        // 验证生成的代码包含 typed_query! 块的完整结构
378        assert!(output.contains("typed_query! {"));
379        assert!(output.contains("table typed_validate_test {"));
380        assert!(output.contains("id: i64,"));
381        assert!(output.contains("name: String,"));
382        // 块应正确闭合
383        let count_open = output.matches("typed_query! {").count();
384        let count_close = output.matches("}\n}").count();
385        assert_eq!(count_open, 1);
386        assert_eq!(count_close, 1);
387    }
388}
sz_orm_core/schema_gen.rs

sz_orm_core/
schema_gen.rs
