sea-orm-generator/src/generator/models/mod.rs

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<u16>,
    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<u16>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub enable: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub table_name: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub info: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rust_type: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub db_type: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub attributes: Option<bool>,
}
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<Vec<GeneratedFileChunk>> {
    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::<Result<Vec<Table>>>()?;

    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<TableCreateStatement>,
    config: ModelConfig,
    writer_context: EntityWriterContext,
) -> Result<Vec<GeneratedFileChunk>> {
    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::<Vec<_>>())
}