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

use std::path::PathBuf;

use crate::generator::file::pathbuf_to_rust_path;

use super::{
    discovery::{table::Table, DiscoveredSchema},
    sea_orm::SeaOrmConfig,
    templates::TemplateConfig,
    Module, ModulesContext,
};
use color_eyre::{
    eyre::{eyre, Context, ContextCompat},
    Result,
};
use minijinja::Environment;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(default)]
pub struct ModelsConfig {
    pub enable: bool,
    pub path: Option<PathBuf>,
    pub prelude: bool,
}

impl Default for ModelsConfig {
    fn default() -> Self {
        Self {
            enable: false,
            path: None,
            prelude: true,
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct ModelTemplateContext {
    entities_path: Option<String>,
    tables: Option<Vec<Table>>,
    prelude_path: Option<String>,
    table_name: Option<String>,
    config: ModelsConfig,
}

impl ModelTemplateContext {
    pub fn new(
        entities_path: Option<String>,
        prelude_path: Option<String>,
        tables: Option<Vec<Table>>,
        table_name: Option<String>,
        config: ModelsConfig,
    ) -> Self {
        Self {
            entities_path,
            tables,
            prelude_path,
            table_name,
            config,
        }
    }
}

#[derive(Debug)]
pub struct ModelsModule;

#[async_trait::async_trait]
impl Module for ModelsModule {
    fn init(&mut self, ctx: &mut ModulesContext) -> Result<()> {
        ctx.get_config_auto::<ModelsConfig>("modules.model")?;
        Ok(())
    }

    async fn validate(&mut self, ctx: &mut ModulesContext) -> Result<bool> {
        let map = ctx.get_anymap();

        if let (Some(config), Some(template_config), Some(sea_orm_config)) = (
            map.get::<ModelsConfig>(),
            map.get::<TemplateConfig>(),
            map.get::<SeaOrmConfig>(),
        ) {
            if config.enable && !template_config.enable {
                return Err(eyre!(
                    "\"modules.template.enable\" must be enabled to use \"modules.model.enable\""
                ));
            }
            if config.enable && !sea_orm_config.enable {
                return Err(eyre!(
                    "\"modules.sea_orm.enable\" must be enabled to use \"modules.model.enable\""
                ));
            }
            if config.enable && config.path.is_none() {
                return Err(eyre!(
                    "\"modules.model.path\" must be set to use \"modules.model.enable\""
                ));
            }
            Ok(config.enable && template_config.enable)
        } else {
            // One or both keys are missing
            Ok(false)
        }
    }
    async fn execute(&mut self, ctx: &mut ModulesContext) -> Result<()> {
        let mut files: Vec<(PathBuf, String)> = Vec::new();
        let map = ctx.get_anymap();

        if let (Some(config), Some(templates), Some(sea_orm_config), Some(schema)) = (
            map.get::<ModelsConfig>(),
            map.get::<Environment<'static>>(),
            map.get::<SeaOrmConfig>(),
            map.get::<DiscoveredSchema>(),
        ) {
            let models_path = config.path.clone().unwrap();
            tracing::info!(?models_path, "Models path");
            let entities_path = sea_orm_config.path.clone().unwrap();
            let mod_path = models_path.join("mod.rs");

            let relative_entities_path = pathdiff::diff_paths(&entities_path, &mod_path)
                .context("Failed to calculate relative path")?;
            let relative_entities_rust_path = pathbuf_to_rust_path(relative_entities_path);
            let context = ModelTemplateContext::new(
                Some(relative_entities_rust_path.clone()),
                None,
                Some(schema.tables.clone()),
                None,
                config.clone(),
            );
            if config.prelude {
                let prelude = templates
                    .get_template("model_prelude")?
                    .render(&context)
                    .context("Failed to render model prelude part")?;
                files.push((models_path.join("prelude.rs"), prelude));
            }

            let mod_content = templates
                .get_template("model_mod")?
                .render(&context)
                .context("Failed to render model mod")?;
            files.push((mod_path.clone(), mod_content));
            for table in &schema.tables {
                tracing::debug!(?table, "Generating model for table");
                let path = models_path.join(format!("{}.rs", table.name));

                let context = ModelTemplateContext::new(
                    Some(relative_entities_rust_path.clone()),
                    if config.prelude {
                        Some("super::prelude".to_string())
                    } else {
                        None
                    },
                    None,
                    Some(table.name.clone()),
                    config.clone(),
                );

                if path.exists() {
                    tracing::debug!(?path, "Model file already exists");
                    continue;
                }

                let content = templates
                    .get_template("model")?
                    .render(&context)
                    .context("Failed to render model")?;
                files.push((path.clone(), content.clone()));
            }
        } else {
            // One or both keys are missing
        }
        tracing::info!(?files, "Generated model files");
        let file_manager = ctx.get_file_manager_mut();
        for (output_path, content) in files {
            file_manager.insert(&output_path, &content, None)?;
        }
        Ok(())
    }
}