use std::fmt::Write as _;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnSchema {
pub name: String,
pub rust_type: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableSchema {
pub name: String,
pub columns: Vec<ColumnSchema>,
}
pub struct SchemaGenerator {
header: String,
emit_use: bool,
emit_model_structs: bool,
model_derives: String,
}
impl Default for SchemaGenerator {
fn default() -> Self {
Self::new()
}
}
impl SchemaGenerator {
pub fn new() -> Self {
Self {
header: format!(
"// Auto-generated by sz-orm-cli generate schema at {}\n\
// DO NOT EDIT MANUALLY — re-run the command to refresh.\n\
//\n\
// This file contains typed_query! table declarations\n\
// enabling compile-time column-name verification.\n",
chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC")
),
emit_use: true,
emit_model_structs: false,
model_derives: "Debug, Clone, serde::Serialize, serde::Deserialize".to_string(),
}
}
pub fn with_header(mut self, header: impl Into<String>) -> Self {
self.header = header.into();
self
}
pub fn emit_use(mut self, emit: bool) -> Self {
self.emit_use = emit;
self
}
pub fn emit_model_structs(mut self, emit: bool) -> Self {
self.emit_model_structs = emit;
self
}
pub fn with_model_derives(mut self, derives: impl Into<String>) -> Self {
self.model_derives = derives.into();
self
}
pub fn generate(&self, tables: &[TableSchema]) -> String {
let mut out = String::new();
writeln!(out, "{}", self.header).expect("write to String is infallible");
writeln!(out).expect("write to String is infallible");
if self.emit_use {
writeln!(out, "use sz_orm_core::typed_query;").expect("write to String is infallible");
if self.emit_model_structs
&& !self.model_derives.is_empty()
&& self.model_derives.contains("serde::")
{
writeln!(out, "use serde::{{Serialize, Deserialize}};")
.expect("write to String is infallible");
}
writeln!(out).expect("write to String is infallible");
}
if self.emit_model_structs {
for (idx, table) in tables.iter().enumerate() {
if idx > 0 {
writeln!(out).expect("write to String is infallible");
}
write!(out, "{}", self.render_model_struct(table))
.expect("write to String is infallible");
}
if !tables.is_empty() {
writeln!(out).expect("write to String is infallible");
writeln!(out).expect("write to String is infallible");
}
}
for (idx, table) in tables.iter().enumerate() {
if idx > 0 {
writeln!(out).expect("write to String is infallible");
}
write!(out, "{}", self.render_table(table)).expect("write to String is infallible");
}
out
}
fn render_table(&self, table: &TableSchema) -> String {
let mut out = String::new();
writeln!(out, "typed_query! {{").expect("write to String is infallible");
writeln!(out, " table {} {{", table.name).expect("write to String is infallible");
for col in &table.columns {
writeln!(out, " {}: {},", col.name, col.rust_type)
.expect("write to String is infallible");
}
writeln!(out, " }}").expect("write to String is infallible");
writeln!(out, "}}").expect("write to String is infallible");
out
}
fn render_model_struct(&self, table: &TableSchema) -> String {
let struct_name = to_pascal_case(&table.name);
let mut out = String::new();
if !self.model_derives.is_empty() {
writeln!(out, "#[derive({})]", self.model_derives)
.expect("write to String is infallible");
}
writeln!(out, "#[table_name = \"{}\"]", table.name).expect("write to String is infallible");
writeln!(out, "pub struct {} {{", struct_name).expect("write to String is infallible");
for col in &table.columns {
writeln!(out, " pub {}: {},", col.name, col.rust_type)
.expect("write to String is infallible");
}
writeln!(out, "}}").expect("write to String is infallible");
out
}
}
fn to_pascal_case(input: &str) -> String {
let mut result = String::with_capacity(input.len());
let mut next_upper = true;
for ch in input.chars() {
if ch == '_' || ch == '-' || ch == ' ' {
next_upper = true;
continue;
}
if next_upper {
result.push(ch.to_ascii_uppercase());
next_upper = false;
} else {
result.push(ch);
}
}
result
}
pub fn sql_type_to_rust(sql_type: &str, nullable: bool) -> String {
let lower = sql_type.to_lowercase();
let after_param_strip = lower.split('(').next().unwrap_or(&lower).trim();
let base_type = after_param_strip
.strip_suffix(" unsigned")
.or_else(|| after_param_strip.strip_suffix(" zerofill"))
.unwrap_or(after_param_strip)
.trim();
let rust = match base_type {
"tinyint" => "i8",
"smallint" | "int2" | "smallserial" => "i16",
"bigint" | "int8" | "bigserial" => "i64",
"int" | "integer" | "int4" | "mediumint" | "serial" => "i32",
"float8" | "double" | "double precision" => "f64",
"float4" | "real" | "float" => "f32",
"decimal" | "numeric" => "f64",
"bool" | "boolean" => "bool",
"blob" | "bytea" | "binary" | "varbinary" | "tinyblob" | "mediumblob" | "longblob" => {
"Vec<u8>"
}
"date" => "String",
"datetime" | "timestamp" | "timestamptz" => "String",
"time" | "timetz" => "String",
"interval" => "String",
"json" | "jsonb" => "String",
"uuid" => "String",
"char" | "varchar" | "text" | "tinytext" | "mediumtext" | "longtext" => "String",
"enum" | "set" => "String",
_ => "String",
};
if nullable {
format!("Option<{}>", rust)
} else {
rust.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sql_type_to_rust_int() {
assert_eq!(sql_type_to_rust("INT", false), "i32");
assert_eq!(sql_type_to_rust("BIGINT", false), "i64");
assert_eq!(sql_type_to_rust("SMALLINT", false), "i16");
assert_eq!(sql_type_to_rust("TINYINT", false), "i8");
}
#[test]
fn test_sql_type_to_rust_float() {
assert_eq!(sql_type_to_rust("FLOAT", false), "f32");
assert_eq!(sql_type_to_rust("DOUBLE", false), "f64");
assert_eq!(sql_type_to_rust("DECIMAL(10,2)", false), "f64");
}
#[test]
fn test_sql_type_to_rust_bool() {
assert_eq!(sql_type_to_rust("BOOLEAN", false), "bool");
assert_eq!(sql_type_to_rust("TINYINT(1)", false), "i8");
}
#[test]
fn test_sql_type_to_rust_string() {
assert_eq!(sql_type_to_rust("VARCHAR(255)", false), "String");
assert_eq!(sql_type_to_rust("TEXT", false), "String");
assert_eq!(sql_type_to_rust("CHAR(36)", false), "String");
}
#[test]
fn test_sql_type_to_rust_binary() {
assert_eq!(sql_type_to_rust("BLOB", false), "Vec<u8>");
assert_eq!(sql_type_to_rust("BYTEA", false), "Vec<u8>");
}
#[test]
fn test_sql_type_to_rust_nullable() {
assert_eq!(sql_type_to_rust("INT", true), "Option<i32>");
assert_eq!(sql_type_to_rust("VARCHAR(255)", true), "Option<String>");
}
#[test]
fn test_sql_type_to_rust_pg_types() {
assert_eq!(sql_type_to_rust("int8", false), "i64");
assert_eq!(sql_type_to_rust("int2", false), "i16");
assert_eq!(sql_type_to_rust("float8", false), "f64");
assert_eq!(sql_type_to_rust("float4", false), "f32");
assert_eq!(sql_type_to_rust("bytea", false), "Vec<u8>");
}
#[test]
fn test_sql_type_to_rust_json_uuid() {
assert_eq!(sql_type_to_rust("JSON", false), "String");
assert_eq!(sql_type_to_rust("JSONB", false), "String");
assert_eq!(sql_type_to_rust("UUID", false), "String");
}
#[test]
fn test_schema_generator_single_table() {
let gen = SchemaGenerator::new().emit_use(false);
let tables = vec![TableSchema {
name: "users".to_string(),
columns: vec![
ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
},
ColumnSchema {
name: "name".to_string(),
rust_type: "String".to_string(),
},
],
}];
let output = gen.generate(&tables);
assert!(output.contains("table users {"));
assert!(output.contains("id: i64,"));
assert!(output.contains("name: String,"));
assert!(output.contains("typed_query! {"));
}
#[test]
fn test_schema_generator_multiple_tables() {
let gen = SchemaGenerator::new().emit_use(false);
let tables = vec![
TableSchema {
name: "users".to_string(),
columns: vec![ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
}],
},
TableSchema {
name: "orders".to_string(),
columns: vec![ColumnSchema {
name: "order_id".to_string(),
rust_type: "i64".to_string(),
}],
},
];
let output = gen.generate(&tables);
assert!(output.contains("table users {"));
assert!(output.contains("table orders {"));
let users_end = output.find("}").unwrap();
let orders_start = output.find("table orders").unwrap();
let between = &output[users_end..orders_start];
assert!(between.contains("\n\n"));
}
#[test]
fn test_schema_generator_with_use_statement() {
let gen = SchemaGenerator::new().emit_use(true);
let tables = vec![TableSchema {
name: "t".to_string(),
columns: vec![],
}];
let output = gen.generate(&tables);
assert!(output.contains("use sz_orm_core::typed_query;"));
}
#[test]
fn test_schema_generator_header() {
let gen = SchemaGenerator::new();
let output = gen.generate(&[]);
assert!(output.contains("Auto-generated"));
assert!(output.contains("DO NOT EDIT MANUALLY"));
}
#[test]
fn test_schema_generator_custom_header() {
let gen = SchemaGenerator::new().with_header("// Custom header\n");
let output = gen.generate(&[]);
assert!(output.starts_with("// Custom header"));
assert!(!output.contains("Auto-generated"));
}
#[test]
fn test_schema_generator_empty_tables() {
let gen = SchemaGenerator::new().emit_use(false);
let output = gen.generate(&[]);
assert!(output.contains("Auto-generated"));
assert!(!output.contains("typed_query! {"));
}
#[test]
fn test_schema_generator_option_type() {
let gen = SchemaGenerator::new().emit_use(false);
let tables = vec![TableSchema {
name: "products".to_string(),
columns: vec![ColumnSchema {
name: "price".to_string(),
rust_type: "Option<f64>".to_string(),
}],
}];
let output = gen.generate(&tables);
assert!(output.contains("price: Option<f64>,"));
}
#[test]
fn test_schema_generator_compound_type() {
let gen = SchemaGenerator::new().emit_use(false);
let tables = vec![TableSchema {
name: "files".to_string(),
columns: vec![ColumnSchema {
name: "content".to_string(),
rust_type: "Vec<u8>".to_string(),
}],
}];
let output = gen.generate(&tables);
assert!(output.contains("content: Vec<u8>,"));
}
#[test]
fn test_generated_code_is_valid_syntax() {
let gen = SchemaGenerator::new().emit_use(false);
let tables = vec![TableSchema {
name: "typed_validate_test".to_string(),
columns: vec![
ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
},
ColumnSchema {
name: "name".to_string(),
rust_type: "String".to_string(),
},
],
}];
let output = gen.generate(&tables);
assert!(output.contains("typed_query! {"));
assert!(output.contains("table typed_validate_test {"));
assert!(output.contains("id: i64,"));
assert!(output.contains("name: String,"));
let count_open = output.matches("typed_query! {").count();
let count_close = output.matches("}\n}").count();
assert_eq!(count_open, 1);
assert_eq!(count_close, 1);
}
#[test]
fn test_sql_type_to_rust_interval_not_int() {
assert_eq!(sql_type_to_rust("INTERVAL", false), "String");
assert_eq!(sql_type_to_rust("interval", false), "String");
assert_eq!(sql_type_to_rust("INTERVAL DAY TO SECOND", false), "String");
}
#[test]
fn test_sql_type_to_rust_mediumint() {
assert_eq!(sql_type_to_rust("MEDIUMINT", false), "i32");
assert_eq!(sql_type_to_rust("MEDIUMINT(8)", false), "i32");
}
#[test]
fn test_sql_type_to_rust_integer_alias() {
assert_eq!(sql_type_to_rust("INTEGER", false), "i32");
assert_eq!(sql_type_to_rust("INTEGER(11)", false), "i32");
}
#[test]
fn test_sql_type_to_rust_serial_types() {
assert_eq!(sql_type_to_rust("SERIAL", false), "i32");
assert_eq!(sql_type_to_rust("BIGSERIAL", false), "i64");
assert_eq!(sql_type_to_rust("SMALLSERIAL", false), "i16");
}
#[test]
fn test_sql_type_to_rust_unsigned_suffix() {
assert_eq!(sql_type_to_rust("INT UNSIGNED", false), "i32");
assert_eq!(sql_type_to_rust("BIGINT UNSIGNED", false), "i64");
assert_eq!(sql_type_to_rust("TINYINT UNSIGNED", false), "i8");
assert_eq!(sql_type_to_rust("SMALLINT UNSIGNED", false), "i16");
assert_eq!(sql_type_to_rust("MEDIUMINT UNSIGNED", false), "i32");
}
#[test]
fn test_sql_type_to_rust_zerofill_suffix() {
assert_eq!(sql_type_to_rust("INT ZEROFILL", false), "i32");
assert_eq!(sql_type_to_rust("INT(4) ZEROFILL", false), "i32");
}
#[test]
fn test_sql_type_to_rust_timestamptz() {
assert_eq!(sql_type_to_rust("TIMESTAMPTZ", false), "String");
assert_eq!(sql_type_to_rust("timestamptz", false), "String");
}
#[test]
fn test_sql_type_to_rust_timetz() {
assert_eq!(sql_type_to_rust("TIMETZ", false), "String");
assert_eq!(sql_type_to_rust("timetz", false), "String");
}
#[test]
fn test_sql_type_to_rust_double_precision() {
assert_eq!(sql_type_to_rust("DOUBLE PRECISION", false), "f64");
assert_eq!(sql_type_to_rust("double precision", false), "f64");
}
#[test]
fn test_sql_type_to_rust_varbinary() {
assert_eq!(sql_type_to_rust("VARBINARY(255)", false), "Vec<u8>");
assert_eq!(sql_type_to_rust("VARBINARY", false), "Vec<u8>");
}
#[test]
fn test_sql_type_to_rust_blob_variants() {
assert_eq!(sql_type_to_rust("TINYBLOB", false), "Vec<u8>");
assert_eq!(sql_type_to_rust("MEDIUMBLOB", false), "Vec<u8>");
assert_eq!(sql_type_to_rust("LONGBLOB", false), "Vec<u8>");
}
#[test]
fn test_sql_type_to_rust_text_variants() {
assert_eq!(sql_type_to_rust("TINYTEXT", false), "String");
assert_eq!(sql_type_to_rust("MEDIUMTEXT", false), "String");
assert_eq!(sql_type_to_rust("LONGTEXT", false), "String");
}
#[test]
fn test_sql_type_to_rust_enum_and_set() {
assert_eq!(sql_type_to_rust("ENUM('a','b')", false), "String");
assert_eq!(sql_type_to_rust("SET('a','b')", false), "String");
}
#[test]
fn test_sql_type_to_rust_decimal_with_space() {
assert_eq!(sql_type_to_rust("DECIMAL(10, 2)", false), "f64");
assert_eq!(sql_type_to_rust("NUMERIC(8, 2)", false), "f64");
}
#[test]
fn test_sql_type_to_rust_unknown_defaults_string() {
assert_eq!(sql_type_to_rust("UNKNOWN_TYPE", false), "String");
assert_eq!(sql_type_to_rust("citext", false), "String");
assert_eq!(sql_type_to_rust("money", false), "String");
}
#[test]
fn test_sql_type_to_rust_interval_nullable() {
assert_eq!(sql_type_to_rust("INTERVAL", true), "Option<String>");
}
#[test]
fn test_to_pascal_case_basic() {
assert_eq!(to_pascal_case("users"), "Users");
assert_eq!(to_pascal_case("orders"), "Orders");
}
#[test]
fn test_to_pascal_case_snake_case() {
assert_eq!(to_pascal_case("order_items"), "OrderItems");
assert_eq!(to_pascal_case("user_profiles"), "UserProfiles");
}
#[test]
fn test_to_pascal_case_kebab_case() {
assert_eq!(to_pascal_case("user-profile"), "UserProfile");
assert_eq!(to_pascal_case("order-item"), "OrderItem");
}
#[test]
fn test_to_pascal_case_with_spaces() {
assert_eq!(to_pascal_case("user profile"), "UserProfile");
}
#[test]
fn test_to_pascal_case_single_char() {
assert_eq!(to_pascal_case("a"), "A");
assert_eq!(to_pascal_case("_a"), "A");
}
#[test]
fn test_to_pascal_case_already_pascal() {
assert_eq!(to_pascal_case("Users"), "Users");
}
#[test]
fn test_emit_model_structs_generates_struct() {
let gen = SchemaGenerator::new()
.emit_use(false)
.emit_model_structs(true);
let tables = vec![TableSchema {
name: "users".to_string(),
columns: vec![
ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
},
ColumnSchema {
name: "name".to_string(),
rust_type: "String".to_string(),
},
],
}];
let output = gen.generate(&tables);
assert!(output.contains("pub struct Users {"));
assert!(output.contains("pub id: i64,"));
assert!(output.contains("pub name: String,"));
assert!(output.contains("#[derive("));
assert!(output.contains("Debug"));
assert!(output.contains("Clone"));
assert!(output.contains("#[table_name = \"users\"]"));
assert!(output.contains("typed_query! {"));
assert!(output.contains("table users {"));
}
#[test]
fn test_emit_model_structs_disabled_by_default() {
let gen = SchemaGenerator::new().emit_use(false);
let tables = vec![TableSchema {
name: "users".to_string(),
columns: vec![ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
}],
}];
let output = gen.generate(&tables);
assert!(!output.contains("pub struct Users"));
assert!(output.contains("typed_query! {"));
}
#[test]
fn test_emit_model_structs_multiple_tables() {
let gen = SchemaGenerator::new()
.emit_use(false)
.emit_model_structs(true);
let tables = vec![
TableSchema {
name: "users".to_string(),
columns: vec![ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
}],
},
TableSchema {
name: "order_items".to_string(),
columns: vec![ColumnSchema {
name: "item_id".to_string(),
rust_type: "i64".to_string(),
}],
},
];
let output = gen.generate(&tables);
assert!(output.contains("pub struct Users {"));
assert!(output.contains("pub struct OrderItems {"));
}
#[test]
fn test_emit_model_structs_with_custom_derives() {
let gen = SchemaGenerator::new()
.emit_use(false)
.emit_model_structs(true)
.with_model_derives("Debug, Clone");
let tables = vec![TableSchema {
name: "users".to_string(),
columns: vec![ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
}],
}];
let output = gen.generate(&tables);
assert!(output.contains("#[derive(Debug, Clone)]"));
assert!(!output.contains("serde::Serialize"));
}
#[test]
fn test_emit_model_structs_with_empty_derives() {
let gen = SchemaGenerator::new()
.emit_use(false)
.emit_model_structs(true)
.with_model_derives("");
let tables = vec![TableSchema {
name: "users".to_string(),
columns: vec![ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
}],
}];
let output = gen.generate(&tables);
assert!(!output.contains("#[derive("));
assert!(output.contains("pub struct Users {"));
}
#[test]
fn test_emit_model_structs_with_nullable_fields() {
let gen = SchemaGenerator::new()
.emit_use(false)
.emit_model_structs(true);
let tables = vec![TableSchema {
name: "products".to_string(),
columns: vec![
ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
},
ColumnSchema {
name: "price".to_string(),
rust_type: "Option<f64>".to_string(),
},
ColumnSchema {
name: "description".to_string(),
rust_type: "Option<String>".to_string(),
},
],
}];
let output = gen.generate(&tables);
assert!(output.contains("pub struct Products {"));
assert!(output.contains("pub id: i64,"));
assert!(output.contains("pub price: Option<f64>,"));
assert!(output.contains("pub description: Option<String>,"));
}
#[test]
fn test_emit_model_structs_with_serde_use() {
let gen = SchemaGenerator::new()
.emit_use(true)
.emit_model_structs(true);
let tables = vec![TableSchema {
name: "users".to_string(),
columns: vec![ColumnSchema {
name: "id".to_string(),
rust_type: "i64".to_string(),
}],
}];
let output = gen.generate(&tables);
assert!(output.contains("use serde::{Serialize, Deserialize};"));
}
#[test]
fn test_emit_model_structs_empty_tables() {
let gen = SchemaGenerator::new()
.emit_use(false)
.emit_model_structs(true);
let output = gen.generate(&[]);
assert!(!output.contains("pub struct"));
assert!(output.contains("Auto-generated"));
}
}