use crate::config::{sea_orm_config::DateTimeCrate, Config}; use color_eyre::Result; use discover::DbType; use file::FileGenerator; use handlebars::Handlebars; use sea_orm_codegen::{EntityTransformer, EntityWriterContext, OutputFile}; use sea_schema::sea_query::TableCreateStatement; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use table::Table; use super::file::{GeneratedFile, GeneratedFileChunk}; pub mod column; pub mod comment; pub mod discover; pub mod file; pub mod table; #[derive(Debug, Clone)] pub struct ModelConfig { pub models_path: PathBuf, pub prelude: bool, pub entities_path: PathBuf, pub relative_entities_path: String, pub enable: bool, pub comment: CommentConfig, pub db_type: DbType, } #[derive(Debug, Clone)] pub struct CommentConfig { pub max_width: Option, pub enable: bool, pub table_name: bool, pub column_info: bool, pub column_name: bool, pub column_rust_type: bool, pub column_db_type: bool, pub column_attributes: bool, pub ignore_errors: bool, pub date_time_crate: DateTimeCrate, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CommentConfigSerde { #[serde(skip_serializing_if = "Option::is_none")] pub max_width: Option, #[serde(skip_serializing_if = "Option::is_none")] pub enable: Option, #[serde(skip_serializing_if = "Option::is_none")] pub table_name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub info: Option, #[serde(skip_serializing_if = "Option::is_none")] pub rust_type: Option, #[serde(skip_serializing_if = "Option::is_none")] pub db_type: Option, #[serde(skip_serializing_if = "Option::is_none")] pub attributes: Option, } impl CommentConfigSerde { pub fn merge(&self, config: &CommentConfig) -> CommentConfig { CommentConfig { max_width: self.max_width.or(config.max_width), table_name: self.table_name.unwrap_or(config.table_name), column_name: self.name.unwrap_or(config.column_name), column_info: self.info.unwrap_or(config.column_info), column_rust_type: self.rust_type.unwrap_or(config.column_rust_type), column_db_type: self.db_type.unwrap_or(config.column_db_type), column_attributes: self.attributes.unwrap_or(config.column_attributes), ignore_errors: config.ignore_errors, enable: self.enable.unwrap_or(config.enable), date_time_crate: config.date_time_crate.clone(), } } } impl ModelConfig { pub fn new(config: Config, db_type: DbType) -> Self { let models_path = config.output.path.join(&config.output.models.path); let entities_path = models_path.join(&config.output.models.entities); ModelConfig { db_type, prelude: config.output.models.prelude, entities_path, models_path, relative_entities_path: config.output.models.entities.clone(), enable: config.output.models.enable, comment: CommentConfig { max_width: config.output.models.comment.max_width, enable: config.output.models.comment.enable, table_name: config.output.models.comment.table_name, column_name: config.output.models.comment.column_name, column_info: config.output.models.comment.column_info, column_rust_type: config.output.models.comment.column_rust_type, column_db_type: config.output.models.comment.column_db_type, column_attributes: config.output.models.comment.column_attributes, ignore_errors: config.output.models.comment.ignore_errors, date_time_crate: config.sea_orm.entity.date_time_crate, }, } } } pub async fn generate_models<'a>( database_url: &str, config: &Config, handlebars: &'a Handlebars<'a>, ) -> Result> { let mut files = Vec::new(); let db_filter = config.sea_orm.entity.tables.get_filter(); let (table_stmts, db_type) = discover::get_tables(database_url.to_owned(), db_filter, &config.db).await?; let model_config = ModelConfig::new(config.clone(), db_type); let writer_context = config.clone().into(); files.extend( generate_entities(table_stmts.clone(), model_config.clone(), writer_context).await?, ); files.push(GeneratedFileChunk { path: model_config.models_path.join("mod.rs"), content: format!("pub mod {};", model_config.relative_entities_path), priority: 0, }); let tables = table_stmts .into_iter() .map(Table::new) .collect::>>()?; if model_config.enable { for table in tables { files.extend(FileGenerator::generate_file(table, &model_config, handlebars).await?); } } Ok(files) } pub async fn generate_entities( table_statements: Vec, config: ModelConfig, writer_context: EntityWriterContext, ) -> Result> { let output = EntityTransformer::transform(table_statements)?.generate(&writer_context); Ok(output .files .into_iter() .map(|OutputFile { name, content }| GeneratedFileChunk { path: config.entities_path.join(name), content, priority: 0, }) .collect::>()) }